Merge pull request #2 from viartemev/orchestrator

RAG implementation

Author: viartemev

README.md

@@ -5,11 +5,18 @@
 import os
 import json
 import logging
+import re
 from openai import OpenAI, APIError
 from openai.types.chat import ChatCompletionMessage
 from dotenv import load_dotenv
 import inspect
 
+from app.rag.retriever import KnowledgeRetriever
+from app.agents.tools import (
+    read_file_tool, read_file_tool_def,
+    list_files_tool, list_files_tool_def,
+)
+
 # Новый, улучшенный системный промпт, превращающий агента в программиста.
 SYSTEM_PROMPT = """
 Ты — автономный AI-агент, способный выполнять сложные задачи, используя доступные инструменты.
@@ -55,6 +62,8 @@ def __init__(
         api_key: str,
         model: str = "o4-mini",
         max_iterations: int = 10,
+        use_rag: bool = False,
+        rag_config: Optional[Dict[str, Any]] = None,
     ):
         self.name = name
         self.role = role
@@ -67,6 +76,25 @@ def __init__(
         self.max_iterations = max_iterations
         self.system_prompt = "Ты — универсальный AI-ассистент."
 
+        # Add default tools
+        self.add_tool(read_file_tool, read_file_tool_def)
+        self.add_tool(list_files_tool, list_files_tool_def)
+        
+        # RAG specific attributes
+        self.use_rag = use_rag
+        self.rag_config = rag_config or {}
+        self.retriever: Optional[KnowledgeRetriever] = None
+        if self.use_rag:
+            try:
+                self.retriever = KnowledgeRetriever()
+            except FileNotFoundError as e:
+                logging.warning(
+                    f"Agent '{self.name}' was configured to use RAG, "
+                    f"but the knowledge base is not found. RAG will be disabled. "
+                    f"Error: {e}"
+                )
+                self.use_rag = False
+
     def add_tool(self, tool_func: Callable, tool_definition: Dict[str, Any]):
         """Добавляет инструмент и его определение."""
         tool_name = tool_func.__name__
@@ -79,6 +107,58 @@ def get_openai_tools(self) -> Optional[List[Dict[str, Any]]]:
             return None
         return self.tool_definitions
 
+    def _create_rag_query(self, briefing: str) -> str:
+        """
+        Creates a focused query for RAG from the detailed briefing.
+        Extracts the current task and the last result from history.
+        """
+        task_match = re.search(r"\*\*YOUR CURRENT TASK \(Step .*\):\*\*\n\n\*\*Task:\*\* (.*)\n\*\*Description:\*\* (.*)", briefing)
+        history_match = re.search(r"\*\*EXECUTION HISTORY:\*\*\n(.*)", briefing, re.DOTALL)
+
+        if not task_match:
+            return briefing # Fallback to full briefing
+
+        task = task_match.group(1)
+        description = task_match.group(2)
+        
+        query_parts = [f"Task: {task}", f"Description: {description}"]
+
+        if history_match:
+            history_str = history_match.group(1).strip()
+            # Get the last entry from the history
+            last_entry = history_str.split("- **Step")[0].strip()
+            if last_entry:
+                query_parts.append(f"Context from previous step: {last_entry}")
+        
+        focused_query = "\n".join(query_parts)
+        logging.info(f"Created focused RAG query for {self.name}: '{focused_query}'")
+        return focused_query
+
+    def _enrich_with_knowledge(self, query: str) -> str:
+        """Enriches a query with context from the knowledge base if RAG is enabled."""
+        if not self.use_rag or not self.retriever:
+            return ""
+
+        top_k = self.rag_config.get("top_k", 3)
+        filters = self.rag_config.get("filters", None)
+
+        logging.info(f"Agent '{self.name}' is retrieving knowledge with top_k={top_k}, filters={filters}")
+        retrieved_knowledge = self.retriever.retrieve(query=query, top_k=top_k, filters=filters)
+
+        if not retrieved_knowledge:
+            logging.info("No specific internal standards found for this query.")
+            return ""
+
+        formatted_knowledge = "\n\n---\n\n".join(
+            [f"Source: {chunk['source']}\n\n{chunk['text']}" for chunk in retrieved_knowledge]
+        )
+        knowledge_context = (
+            "Before you begin, consult these internal standards and best practices:"
+            f"\n\n--- RELEVANT KNOWLEDGE ---\n{formatted_knowledge}\n--------------------------\n"
+        )
+        logging.info(f"Knowledge context added for agent '{self.name}'.")
+        return knowledge_context
+
     def _execute_tool(self, tool_name: str, tool_args: Dict[str, Any]) -> Any:
         """Выполняет указанный инструмент с аргументами."""
         if tool_name in self.tools:
@@ -126,9 +206,17 @@ def execute_task(self, briefing: str) -> str:
         """
         logging.info(f"Агент {self.name} получил задачу.")
 
+        knowledge_context = ""
+        if self.use_rag:
+            focused_query = self._create_rag_query(briefing)
+            knowledge_context = self._enrich_with_knowledge(query=focused_query)
+        
         # Системный промпт определяет "личность" агента, а брифинг - контекст задачи.
+        # Контекст из базы знаний добавляется в начало системного промпта.
+        final_system_prompt = f"{knowledge_context}\n{self.system_prompt}"
+
         self.conversation_history = [
-            {"role": "system", "content": self.system_prompt},
+            {"role": "system", "content": final_system_prompt},
             {"role": "user", "content": briefing}
         ]
 

app/agents/agent.py

@@ -1,30 +1,59 @@
 """Coding Agent."""
 from typing import List, Dict
 from app.agents.agent import Agent
+from app.agents.tools import (
+    read_file_tool, read_file_tool_def,
+    edit_file_tool, edit_file_tool_def,
+    list_files_tool, list_files_tool_def,
+)
 
 
 class CodingAgent(Agent):
     """An agent specializing in writing and refactoring code."""
 
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, name: str, role: str, goal: str, **kwargs):
+        super().__init__(
+            name=name,
+            role=role,
+            goal=goal,
+            **kwargs,
+        )
+        # Self-register tools
+        self.add_tool(read_file_tool, read_file_tool_def)
+        self.add_tool(edit_file_tool, edit_file_tool_def)
+        self.add_tool(list_files_tool, list_files_tool_def)
+
         self.system_prompt = (
             "You are a CodingAgent, an elite AI developer. Your task is to write, modify, and fix code."
             "You will be provided with the full plan, the history of previous steps, and your current task."
             "\n\n"
             "## Operating Principles:\n"
             "1.  **Think Before You Code:** Carefully study the task and context. Plan your actions."
-            "2.  **Follow Instructions:** Precisely follow the given task, whether it's writing a new function, fixing a bug, or refactoring."
-            "3.  **Use Tools Wisely:** Do not call tools unnecessarily. Analyze first, then act."
+            "2.  **Surgical Edits:** Your primary goal is to make precise, targeted changes. Do not rewrite entire files. Instead, identify the specific function, method, or block of code that needs changing and replace only that part.\n"
+            "3.  **Use Tools Wisely:**\n"
+            "    - Prefer `edit_file_tool` with `mode='replace'`. This is the safest and most professional way to work.\n"
+            "    - Use `mode='append'` for adding new functions or tests to the end of a file.\n"
+            "    - Use `mode='overwrite'` ONLY when creating a brand new file.\n"
             "4.  **Code Quality:** Write clean, efficient, and well-documented code that adheres to PEP8."
-            "5.  **Handling Code Review Feedback:** "
-            "   - Carefully review ALL feedback from the ReviewerAgent."
-            "   - **'Read-Modify-Overwrite' Strategy:** Instead of many small fixes, use the following approach:"
-            "     a. Read the file's content (`read_file_tool`)."
-            "     b. Apply ALL necessary changes in memory."
-            "     c. Completely overwrite the file with a single call to `edit_file_tool` using `mode='overwrite'` and the full new content."
-            "   - This approach ensures that all corrections are applied atomically and nothing is missed."
             "\n\n"
+            "## Example of a Surgical Edit:\n"
+            "Your task is to fix a bug in the `add` function in `math_utils.py`.\n\n"
+            "1.  **First, read the file:** `read_file_tool(path='app/utils/math_utils.py')`\n"
+            "2.  **Identify the flawed function:**\n"
+            "    ```python\n"
+            "    # This is the old, incorrect code block you will replace\n"
+            "    def add(a, b):\n"
+            "        return a - b # Bug is here\n"
+            "    ```\n"
+            "3.  **Call the edit tool to replace ONLY that function:**\n"
+            "    ```python\n"
+            "    edit_file_tool(\n"
+            "        path='app/utils/math_utils.py',\n"
+            "        mode='replace',\n"
+            "        old_content='''def add(a, b):\\n    return a - b # Bug is here''',\n"
+            "        new_content='''def add(a, b):\\n    # Fix: Correctly perform addition\\n    return a + b'''\n"
+            "    )\n"
+            "    ```\n\n"
             "Your goal is to successfully complete your part of the plan, preparing the way for the next agent."
         )
 

app/agents/roles/coding_agent.py

@@ -1,10 +1,21 @@
 """Evaluator Agent."""
 from app.agents.agent import Agent
+from app.agents.tools import (
+    read_file_tool, read_file_tool_def,
+)
 
 class EvaluatorAgent(Agent):
     """An agent that analyzes errors and results."""
-    def __init__(self, **kwargs):
-        super().__init__(**kwargs)
+    def __init__(self, name: str, role: str, goal: str, **kwargs):
+        super().__init__(
+            name=name,
+            role=role,
+            goal=goal,
+            **kwargs,
+        )
+        # Self-register tools
+        self.add_tool(read_file_tool, read_file_tool_def)
+        
         self.system_prompt = """
 You are EvaluatorAgent, an experienced QA engineer and systems analyst.
 Your primary task is to analyze the log from failed pytest runs and formulate a clear, concise, and single task for the CodingAgent to fix the code.

app/agents/roles/evaluator_agent.py

@@ -1,45 +1,25 @@
 """Reviewer Agent."""
 from app.agents.agent import Agent
-from app.rag.retriever import KnowledgeRetriever
+from app.agents.tools import (
+    read_file_tool, read_file_tool_def,
+)
 
 class ReviewerAgent(Agent):
     """An agent specializing in strict Code Review."""
-    def __init__(self, name: str = "CodeReviewer", **kwargs):
+    def __init__(self, name: str, role: str, goal: str, **kwargs):
         super().__init__(
             name=name,
-            role="Code Reviewer",
-            goal=(
-                "Ensure that the provided code is of high quality, "
-                "free of errors, and adheres to best practices and "
-                "internal coding standards."
-            ),
+            role=role,
+            goal=goal,
             **kwargs,
         )
-        self.retriever = KnowledgeRetriever()
+        # Self-register tools
+        self.add_tool(read_file_tool, read_file_tool_def)
 
-    def _get_system_prompt(self, **kwargs) -> str:
-        code_to_review = kwargs.get("code", "")
-
-        # Retrieve relevant knowledge
-        retrieved_knowledge = self.retriever.retrieve(query=code_to_review, top_k=3)
-        
-        knowledge_context = "No specific internal standards found for this code."
-        if retrieved_knowledge:
-            formatted_knowledge = "\n\n---\n\n".join(
-                [f"Source: {chunk['source']}\n\n{chunk['text']}" for chunk in retrieved_knowledge]
-            )
-            knowledge_context = (
-                "When performing the review, pay close attention to the following "
-                "internal standards and best practices:\n\n"
-                f"--- RELEVANT KNOWLEDGE ---\n{formatted_knowledge}\n--------------------------"
-            )
-
-        return (
-            f"You are a Senior Software Engineer acting as a code reviewer. "
-            f"Your task is to provide a thorough review of the given code snippet.\n\n"
-            f"{knowledge_context}\n\n"
-            f"Please review the following code:\n\n"
-            f"```python\n{code_to_review}\n```\n\n"
-            f"Provide your feedback in a clear, constructive manner. "
-            f"If you find issues, suggest specific improvements."
+        self.system_prompt = (
+            "You are a Senior Software Engineer acting as a code reviewer. "
+            "Your task is to provide a thorough review of the code based on the "
+            "provided file path. Use your available tools to read the file content.\n\n"
+            "If the code meets all standards, respond with only the word 'LGTM'.\n"
+            "Otherwise, provide clear, constructive feedback on what needs to be improved."
         )

app/agents/roles/reviewer_agent.py

@@ -0,0 +1,97 @@
+"""Task Decomposer Agent."""
+from app.agents.agent import Agent
+import logging
+import json
+
+class TaskDecomposer(Agent):
+    """
+    An agent specializing in breaking down a high-level goal into a step-by-step plan.
+    """
+    def __init__(self, name: str, role: str, goal: str, **kwargs):
+        # This agent typically does not need RAG, but the option is there.
+        kwargs.setdefault('use_rag', False)
+        
+        super().__init__(
+            name=name,
+            role=role,
+            goal=goal,
+            **kwargs,
+        )
+        self.system_prompt = """
+You are an expert project manager. Your task is to break down a high-level user goal into a concise, step-by-step plan.
+Each step must be a single, clear action assigned to one of the available roles.
+
+# AVAILABLE ROLES:
+- CodingAgent: Writes, modifies, and fixes code.
+- TestingAgent: Runs tests and reports results.
+- ReviewerAgent: Performs code reviews, checking for quality and adherence to standards.
+- EvaluatorAgent: Analyzes test failures and creates bug reports.
+
+# OUTPUT FORMAT:
+Your output must be a list of steps in JSON format. Do not include any other text or explanation.
+
+# EXAMPLE:
+Goal: "Create a function to add two numbers and test it."
+
+Your output:
+```json
+[
+  {
+    "step": 1,
+    "assignee": "CodingAgent",
+    "task": "Create a new function 'add(a, b)' in 'app/utils/math.py'."
+  },
+  {
+    "step": 2,
+    "assignee": "ReviewerAgent",
+    "task": "Review the 'add' function in 'app/utils/math.py'."
+  },
+  {
+    "step": 3,
+    "assignee": "CodingAgent",
+    "task": "Create a new test file 'tests/test_math.py' with tests for the 'add' function."
+  },
+  {
+    "step": 4,
+    "assignee": "TestingAgent",
+    "task": "Run the tests in 'tests/test_math.py'."
+  }
+]
+```
+"""
+
+    def get_plan(self, goal: str) -> list:
+        """
+        Generates a plan for a given goal.
+        Overrides the base 'execute_task' to return a structured plan.
+        """
+        task_briefing = f"Create a step-by-step plan to achieve the following goal: {goal}"
+        
+        try:
+            logging.info("Requesting plan from OpenAI...")
+            response = self.client.chat.completions.create(
+                model=self.model,
+                messages=[
+                    {"role": "system", "content": self.system_prompt},
+                    {"role": "user", "content": task_briefing},
+                ],
+                response_format={"type": "json_object"},
+            )
+            
+            response_content = response.choices[0].message.content
+            logging.info("Received raw plan: %s", response_content)
+            
+            # The response is a JSON string, so we need to parse it.
+            plan = json.loads(response_content)
+            # Sometimes the model returns a dictionary with a "plan" key
+            if isinstance(plan, dict) and "plan" in plan:
+                return plan["plan"]
+            return plan
+
+        except json.JSONDecodeError as e:
+            logging.error(f"Failed to decode JSON from OpenAI response: {e}")
+            logging.error(f"Raw response was: {response_content}")
+            return [{"step": 1, "assignee": "DefaultAgent", "task": "Failed to create a valid plan due to JSON error."}]
+        except Exception as e:
+            logging.error(f"An unexpected error occurred while getting the plan: {e}", exc_info=True)
+            return [{"step": 1, "assignee": "DefaultAgent", "task": "Failed to create a plan due to an unexpected error."}]