added milvus vector database for RAG operations.

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)

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,22 +28,27 @@
 
 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)
         tester_agent = agents.tester_agent(tester_tools)
         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,16 +62,21 @@ 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 #######")
     print("---------------------------------")
     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)

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

requirements.txt

@@ -2,3 +2,5 @@ crewai
 crewai_tools
 langchain_openai
 python-decouple
+milvus
+pymilvus

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(

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}"

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."