diff --git a/README.md b/README.md index bbb530e..ab0a8eb 100644 --- a/README.md +++ b/README.md @@ -23,6 +23,30 @@ - `python-decouple` library - `crewai` library - OpenAI API Key +- Milvus vector database + + +## Using Vector Database + +This project uses Milvus as a vector database to store and search for relevant code snippets based on the user's query. The `MilvusClient` class handles the connection and operations with the Milvus server. + +### Setting up Milvus + +1. Install Milvus: + + ```sh + pip install pymilvus + ``` + +2. Start Milvus server: + Follow the [Milvus installation guide](https://milvus.io/docs/install_standalone-docker.md) to set up and start the Milvus server. + +3. Configure the connection: + The `MilvusClient` class in `milvus_client.py` connects to the Milvus server running on `localhost` at port `19530` by default. You can change these settings in the `MilvusClient` class constructor. + +### Example Usage + +The `CustomCrew` class in `main.py` uses the `MilvusClient` to query for relevant code snippets based on the user's input. The retrieved code snippets are then used by various agents to perform their tasks. ## Installation @@ -104,6 +128,7 @@ devain/ - [ ] Add Agent Logs - [ ] Stackoverflow tool - [ ] Code Execution +- [ ] Develop on existing codebase instead of generating new code every time [![Star History Chart](https://api.star-history.com/svg?repos=theyashwanthsai/Devyan&type=Date)](https://star-history.com/#theyashwanthsai/Devyan&Date) diff --git a/main.py b/main.py index 56fd80f..ca379db 100644 --- a/main.py +++ b/main.py @@ -5,11 +5,12 @@ from crewai import Agent, Task, Crew, Process from langchain_openai import ChatOpenAI from decouple import config +from milvus_client import MilvusClient from textwrap import dedent from agents import CustomAgents from tasks import CustomTasks -from crewai_tools import FileReadTool +from crewai_tools import FileReadTool from tools.file_write import FileWriteTool from tools.directory_write import DirWriteTool from langchain_community.tools import DuckDuckGoSearchRun @@ -27,14 +28,19 @@ os.environ["OPENAI_API_KEY"] = config("OPENAI_API_KEY") - class CustomCrew: def __init__(self, user_input): self.user_input = user_input + self.milvus_client = MilvusClient() + def run(self): agents = CustomAgents() tasks = CustomTasks() + # Query Milvus for relevant code snippets + query_embedding = self.get_query_embedding(self.user_input) + code_snippets = self.milvus_client.search(query_embedding) + # Agents architect_agent = agents.architect_agent(architect_tools) programmer_agent = agents.programmer_agent(programmer_tools) @@ -42,7 +48,7 @@ def run(self): reviewer_agent = agents.reviewer_agent(reviewer_tools) # Tasks - architecture_task = tasks.architecture_task(architect_agent, architect_tools, self.user_input) + architecture_task = tasks.architecture_task(architect_agent, architect_tools, self.user_input, code_snippets) implementation_task = tasks.implementation_task(programmer_agent, programmer_tools, [architecture_task]) testing_task = tasks.testing_task(tester_agent, tester_tools, [implementation_task]) reviewing_task = tasks.reviewing_task(reviewer_agent, reviewer_tools, [architecture_task, implementation_task, testing_task]) @@ -56,7 +62,12 @@ def run(self): result = crew.kickoff() return result - + def get_query_embedding(self, query): + # Convert the query to an embedding using a pre-trained model + from langchain.embeddings import OpenAIEmbeddings + embeddings = OpenAIEmbeddings() + query_embedding = embeddings.embed_query(query) + return query_embedding if __name__ == "__main__": print("\n####### Welcome to Devyan #######") @@ -64,8 +75,8 @@ def run(self): user_input = input("What problem do you want me to solve?\n") crew = CustomCrew(user_input) result = crew.run() - + print("\n\n########################") - print("## Here is you crew run result:") + print("## Here is your crew run result:") print("########################\n") - print(result) + print(result) \ No newline at end of file diff --git a/milvus_client.py b/milvus_client.py new file mode 100644 index 0000000..81c5205 --- /dev/null +++ b/milvus_client.py @@ -0,0 +1,33 @@ +from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType + + +class MilvusClient: + def __init__(self, host="localhost", port="19530"): + self.host = host + self.port = port + self.collection_name = "codebase" + self.connect() + + def connect(self): + try: + connections.connect(host=self.host, port=self.port) + except Exception as e: + print(f"Failed to connect to Milvus: {e}") + + def create_collection(self): + fields = [ + FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True), + FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768), + ] + schema = CollectionSchema(fields, "Codebase collection") + self.collection = Collection(self.collection_name, schema) + + def insert_embeddings(self, embeddings): + self.collection.insert([embeddings]) + + def search(self, query_embedding, top_k=5): + search_params = {"metric_type": "L2", "params": {"nprobe": 10}} + results = self.collection.search( + [query_embedding], "embedding", search_params, limit=top_k + ) + return results diff --git a/requirements.txt b/requirements.txt index 9914b6c..bfcc575 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,3 +2,5 @@ crewai crewai_tools langchain_openai python-decouple +milvus +pymilvus \ No newline at end of file diff --git a/tasks.py b/tasks.py index a025169..cde85b6 100644 --- a/tasks.py +++ b/tasks.py @@ -5,16 +5,18 @@ class CustomTasks: def __tip_section(self): return "If you do your BEST WORK, I'll give you a $10,000 commission!" - - def architecture_task(self, agent, tools, user_input): + def architecture_task(self, agent, tools, user_input, code_snippets): + code_context = "\n".join([snippet for snippet in code_snippets]) return Task( description=dedent( f""" - Provide a high-level solution architecture for the given problem: {user_input}. - Your final answer must include a clear overview and major components involved. - {self.__tip_section()} - You have access to tools which can search the internet, read files, write files and create directories - """ + Provide a high-level solution architecture for the given problem: {user_input}. + Your final answer must include a clear overview and major components involved. + {self.__tip_section()} + You have access to tools which can search the internet, read files, write files and create directories. + Here are some relevant code snippets from the existing codebase: + {code_context} + """ ), expected_output='A document outlining the high-level architecture.', tools=tools, @@ -25,18 +27,18 @@ def implementation_task(self, agent, tools, context): return Task( description=dedent( f""" - Implement the solution as per the architect's overview. - Your final answer must include the code implementing the solution. - {self.__tip_section()} - You have access to tools which can read files, write files and create directories - """ + Implement the solution as per the architect's overview. + Your final answer must include the code implementing the solution. + {self.__tip_section()} + You have access to tools which can read files, write files and create directories. + """ ), expected_output='Python code (py files) implementing the solution.', tools=tools, agent=agent, context=context ) - + def testing_task(self, agent, tools, context): return Task( description=dedent( diff --git a/tools/directory_write.py b/tools/directory_write.py index 1c9ce15..4590e3e 100644 --- a/tools/directory_write.py +++ b/tools/directory_write.py @@ -8,8 +8,11 @@ class DirWriteTool(): @tool("Create directory") def dir_write_tool(directory_path): """Useful to create a directory with the given path.""" - if not os.path.exists(directory_path): - os.makedirs(directory_path) - return f"Directory '{directory_path}' has been created successfully." - else: - return f"Directory '{directory_path}' already exists." + try: + if not os.path.exists(directory_path): + os.makedirs(directory_path) + return f"Directory '{directory_path}' has been created successfully." + else: + return f"Directory '{directory_path}' already exists." + except Exception as e: + return f"Failed to create directory '{directory_path}': {e}" diff --git a/tools/file_write.py b/tools/file_write.py index a108dfd..7e4c888 100644 --- a/tools/file_write.py +++ b/tools/file_write.py @@ -8,6 +8,10 @@ class FileWriteTool(): @tool("Write file") def file_write_tool(filename, content): """Useful to write content to a file with the given filename.""" - with open(filename, 'w') as file: - file.write(content) + try: + with open(filename, 'w') as file: + file.write(content) + return f"File '{filename}' has been written successfully." + except Exception as e: + return f"Failed to write file '{filename}': {e}" return f"File '{filename}' has been written successfully."