Applied several fixes.

Author: srtab

CHANGELOG.md

@@ -7,6 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Fixed
+
+- `reply_reviewer` node of `ReviewAddressorAgent` was not using the correct tool to reply to the reviewer comments. We completely refactored the agent to turn it more reliable and robust.
+- `SearchCodeSnippetsTool` was being called with `repository` parameter even when `repo_id` was being provided, leading to errors. Now we support conditionally add the `repository` parameter to the signature of the tool.
+- Sometimes `Document.id` was being defined as an uuid when retrieving the document from the database, leading to errors..
+
 ## [0.1.3] - 2025-05-20
 
 ### Added

daiv/automation/agents/codebase_chat/agent.py

@@ -39,7 +39,7 @@ def compile(self) -> CompiledGraph:
         return create_react_agent(
             self.get_model(model=settings.MODEL_NAME, temperature=settings.TEMPERATURE),
             state_schema=CodebaseChatAgentState,
-            tools=[SearchCodeSnippetsTool(api_wrapper=index)],
+            tools=[SearchCodeSnippetsTool(api_wrapper=index, all_repositories=True)],
             prompt=ChatPromptTemplate.from_messages([codebase_chat_system, MessagesPlaceholder("messages")]).partial(
                 repositories=index._get_all_repositories(),
                 search_code_snippets_name=SEARCH_CODE_SNIPPETS_NAME,
@@ -60,7 +60,7 @@ async def acompile(self) -> CompiledGraph:
         return create_react_agent(
             self.get_model(model=settings.MODEL_NAME, temperature=settings.TEMPERATURE),
             state_schema=CodebaseChatAgentState,
-            tools=[SearchCodeSnippetsTool(api_wrapper=index)],
+            tools=[SearchCodeSnippetsTool(api_wrapper=index, all_repositories=True)],
             prompt=ChatPromptTemplate.from_messages([codebase_chat_system, MessagesPlaceholder("messages")]).partial(
                 repositories=await sync_to_async(index._get_all_repositories)(),
                 search_code_snippets_name=SEARCH_CODE_SNIPPETS_NAME,

daiv/automation/agents/codebase_chat/prompts.py

@@ -19,7 +19,7 @@
      “Sorry, I can only help with questions about the repositories i have access.”
    - Otherwise, continue.
 
-2. **Analysis** For repository-related queries, extract:
+2. **Analysis** For repository-related queries, extract:
    - Programming languages / frameworks (with a brief in-code example).
    - Key search terms (ranked by relevance, with how each might appear in code).
    - Main concepts or topics (ranked, with a short why-it-matters note).
@@ -29,7 +29,7 @@
 
 <repository_search>
 - Use **`{{ search_code_snippets_name }}`** only when the query pertains to these repositories.
-- Always follow the tool’s schema exactly.
+- Always follow the tool's schema exactly.
 - Search with the keywords you extracted; batch similar searches together.
 </repository_search>
 
@@ -48,13 +48,14 @@
 ```
 *Omit the “References” section if you did not cite code.*
 </crafting_the_reply>
-
+{% if repositories %}
 <repositories_accessible_to_daiv>
 DAIV has access to the following repositories:
 {% for repository in repositories %}
  - {{ repository }}
 {%- endfor %}
 </repositories_accessible_to_daiv>
+{% endif %}
 """,  # noqa: E501
     "jinja2",
 )

daiv/automation/agents/review_addressor/agent.py

@@ -8,13 +8,14 @@
 from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder, SystemMessagePromptTemplate
 from langchain_core.runnables import Runnable, RunnableConfig
 from langgraph.graph import END, StateGraph
-from langgraph.graph.state import CompiledStateGraph
+from langgraph.graph.state import CompiledGraph, CompiledStateGraph
 from langgraph.prebuilt import create_react_agent
 from langgraph.store.base import BaseStore  # noqa: TC002
 from langgraph.types import Command
 
 from automation.agents import BaseAgent
 from automation.agents.plan_and_execute import PlanAndExecuteAgent
+from automation.tools import think
 from automation.tools.toolkits import ReadRepositoryToolkit, WebSearchToolkit
 from codebase.clients import RepoClient
 from codebase.indexes import CodebaseIndex
@@ -24,7 +25,6 @@
 from .prompts import respond_reviewer_system, review_comment_system, review_plan_system_template
 from .schemas import ReviewCommentEvaluation, ReviewCommentInput
 from .state import OverallState, ReplyAgentState
-from .tools import reply_reviewer_tool
 
 logger = logging.getLogger("daiv.agents")
 
@@ -41,6 +41,28 @@ def compile(self) -> Runnable:
         ).with_config({"run_name": "ReviewCommentEvaluator"})
 
 
+class ReplyReviewerAgent(BaseAgent[CompiledGraph]):
+    """
+    Agent to reply to reviewer's comments or questions.
+    """
+
+    def compile(self) -> CompiledGraph:
+        tools = ReadRepositoryToolkit.create_instance().get_tools() + WebSearchToolkit.create_instance().get_tools()
+
+        return create_react_agent(
+            self.get_model(model=settings.REPLY_MODEL_NAME, temperature=settings.REPLY_TEMPERATURE),
+            state_schema=ReplyAgentState,
+            tools=tools + [think],
+            store=self.store,
+            checkpointer=False,
+            prompt=ChatPromptTemplate.from_messages([respond_reviewer_system, MessagesPlaceholder("messages")]).partial(
+                current_date_time=timezone.now().strftime("%d %B, %Y %H:%M")
+            ),
+            name=settings.REPLY_NAME,
+            version="v2",
+        )
+
+
 class ReviewAddressorAgent(BaseAgent[CompiledStateGraph]):
     """
     Agent to address reviews by providing feedback and asking questions.
@@ -138,26 +160,8 @@ def reply_reviewer(
         Returns:
             Command[Literal["__end__"]]: The next step in the workflow.
         """
-        tools = ReadRepositoryToolkit.create_instance().get_tools() + WebSearchToolkit.create_instance().get_tools()
-
-        react_agent = create_react_agent(
-            self.get_model(model=settings.REPLY_MODEL_NAME, temperature=settings.REPLY_TEMPERATURE),
-            state_schema=ReplyAgentState,
-            tools=tools + [reply_reviewer_tool],
-            store=store,
-            checkpointer=False,
-            # FIXME: Add diff hunk referenced file to the prompt to improve the agent's performance
-            prompt=ChatPromptTemplate.from_messages([respond_reviewer_system, MessagesPlaceholder("messages")]),
-            name="reply_reviewer_react_agent",
-            version="v2",
-        )
+        reply_reviewer_agent = ReplyReviewerAgent(store=store).agent
 
-        result = react_agent.invoke({
-            "messages": state["notes"],
-            "diff": state["diff"],
-            "current_date_time": timezone.now().strftime("%d %B, %Y %H:%M"),
-        })
+        result = reply_reviewer_agent.invoke({"messages": state["notes"], "diff": state["diff"]})
 
-        # The reply is updated in the state by the reply_reviewer tool.
-        # There's cases where the tool is not called, so we use the last message content as the reply.
         return Command(goto=END, update={"reply": result["messages"][-1].content})

daiv/automation/agents/review_addressor/conf.py

@@ -11,6 +11,7 @@ class ReviewAddressorSettings(BaseSettings):
     REVIEW_COMMENT_MODEL_NAME: ModelName | str = Field(
         default=ModelName.GPT_4_1_MINI, description="Model name to be used for review assessment."
     )
+    REPLY_NAME: str = Field(default="ReplyReviewer", description="Name of the reply reviewer agent.")
     REPLY_MODEL_NAME: ModelName | str = Field(
         default=ModelName.GPT_4_1, description="Model name to be used for reply to comments or questions."
     )

daiv/automation/agents/review_addressor/prompts.py

@@ -45,59 +45,69 @@
 )
 
 respond_reviewer_system = SystemMessagePromptTemplate.from_template(
-    """You are a senior software developer.
-Your role is to give **insightful, professional, constructive replies** to comments or questions left on a merge-request review.
-
-_Current date & time: {{ current_date_time }}_
-
-## 1 Context you receive
-* **Reviewer's comment / question**
-* **Diff hunk** - the file name(s) and exact line(s) of code to which the reviewer is referring:
-
-<diff_hunk>
-{{ diff }}
-</diff_hunk>
-
-*You may also call tooling that inspects the wider codebase.*
-
-> **Reference rules for ambiguous words**
-> If the reviewer says “this”, “here”, “below”, etc., assume the word refers to the line(s) shown in the diff hunk or the immediately neighbouring content of that file.
-
-## 2 If the comment is vague
-If the reviewer's message is too ambiguous for a grounded reply, **do not analyse**.
-Instead, call the `answer_reviewer` tool **once** in this turn to ask a clarifying question, then stop.
-Resume the normal flow only after clarification is provided in a later turn.
-
-## 3 Analysis block  *(only when the comment is clear)*
-Wrap your deep-dive analysis inside **exactly one** pair of tags:
-
-```xml
-<analysis>
-  - Restate the reviewer's comment in your own words.
-  - Quote the relevant lines from the diff hunk (include the leading +/- markers if present).
-  - Explain how the comment relates to those lines.
-  - Consider wider code-base context (using tools if helpful).
-  - Discuss functionality, performance, maintainability, edge-cases, and possible bugs.
-  - Suggest improvements or alternatives (do **not** change code directly).
-  - Summarise overall impact and prioritise the findings.
-</analysis>
-```
-
-## 4 Final reply to the reviewer
-Call `reply_reviewer` tool immediately **after** your analysis (or directly, when Step 2 triggered clarification).
-
-* Use first-person (“I suggest…”, “I noticed…”).
-* Use the same language as the reviewer.
-* Provide technical explanations, but **do not add meta text** such as “Here is my answer” or “Hope this helps.”
-* Never mention the term “diff hunk” in the reply.
-
-## 5 Tool-usage conventions
-* You may call other code-inspection tools if needed.
-* If you make multiple *independent* tool calls, place them together in one `function_calls` block.
-
----
-
-Follow the steps above to reply to the reviewer's next comment or question.
+    """You are a senior software engineer tasked with writing **accurate, professional replies** to merge-request review comments.
+
+
+────────────────────────────────────────────────────────
+CURRENT DATE-TIME:  {{ current_date_time }}
+
+INCOMING CONTEXT
+  • Reviewer's comment / question
+  • Code excerpt (file name + exact lines):
+
+    <code_diff>
+    {{ diff }}
+    </code_diff>
+
+AVAILABLE TOOLS
+  • web_search
+  • repository_structure
+  • retrieve_file_content
+  • search_code_snippets
+  • think   ← private chain-of-thought
+
+────────────────────────────────────────────────────────
+WORKFLOW
+
+### Step 0 • Decide if clarification is needed
+If the reviewer's message is too vague for a grounded answer:
+
+1. Output **one** clarifying question addressed to the reviewer.
+2. Do **not** call any tools.
+3. End the turn.
+
+### Step 1 • Decide whether extra context is required
+Ask yourself: *“Can I answer confidently from the diff alone?”*
+• **If yes** → skip directly to Step 2.
+• **If no** → call whichever inspection tools supply the missing context.
+  - Group multiple calls in a single JSON array.
+  - Stop once you have enough information.
+
+### Step 2 • Private reasoning
+Call the `think` tool **exactly once**, with a `thought` field that includes:
+  • Why you did or did not need extra tools.
+  • Insights gleaned from any tool responses.
+  • How these insights address the reviewer's comment.
+  • Discussion of functionality, performance, maintainability, edge-cases, bugs.
+  • Suggested improvements (do **not** edit code directly).
+  • Impact / priority summary.
+(≈ 250 words max; this content is never shown to the reviewer.)
+
+### Step 3 • Final reply shown to the reviewer
+Immediately after the `think` call, emit plain text following:
+  • First-person voice (“I suggest…”, “I noticed…”).
+  • Match the reviewer's language if detection is confident; otherwise use English.
+  • Be technically precise, referencing code generically (“the line above/below”); **never** say “diff hunk”.
+  • Concise yet complete; avoid unnecessary verbosity.
+
+────────────────────────────────────────────────────────
+RULES OF THUMB
+• Ground every claim in evidence from the diff or tools; avoid speculation.
+• If you skipped the inspection tools, your `think` notes must state why the diff alone sufficed.
+• Keep total output lean; no superfluous headings or meta comments.
+
+────────────────────────────────────────────────────────
+Follow this workflow for the reviewer's next comment.
 """,  # noqa: E501
     "jinja2",
 )

daiv/automation/tools/repository.py

@@ -2,6 +2,7 @@
 
 import logging
 import textwrap
+from typing import Any
 
 from langchain_core.runnables import RunnableConfig  # noqa: TC002
 from langchain_core.tools import BaseTool
@@ -19,6 +20,7 @@
 
 from .schemas import (
     CreateNewRepositoryFileInput,
+    CrossSearchCodeSnippetsInput,
     DeleteRepositoryFileInput,
     RenameRepositoryFileInput,
     ReplaceSnippetInFileInput,
@@ -48,11 +50,17 @@ class SearchCodeSnippetsTool(BaseTool):
         """  # noqa: E501
     ).format(retrieve_file_content_name=RETRIEVE_FILE_CONTENT_NAME)
 
-    args_schema: type[BaseModel] = SearchCodeSnippetsInput
     handle_validation_error: bool = True
 
     api_wrapper: CodebaseIndex = Field(default_factory=lambda: CodebaseIndex(repo_client=RepoClient.create_instance()))
 
+    def __init__(self, *, all_repositories: bool = False, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        if all_repositories:
+            self.args_schema = CrossSearchCodeSnippetsInput
+        else:
+            self.args_schema = SearchCodeSnippetsInput
+
     def _run(self, query: str, intent: str, config: RunnableConfig, repository: str | None = None) -> str:
         """
         Searches the codebase for a given query.

daiv/automation/tools/schemas.py

@@ -23,14 +23,21 @@ class SearchCodeSnippetsInput(BaseModel):
             """  # noqa: E501
         ),
     )
+    intent: str = Field(..., description="A brief description of why you are searching for this code.")
+
+
+class CrossSearchCodeSnippetsInput(SearchCodeSnippetsInput):
+    """
+    Search for code snippets in specific repositories or all repositories.
+    """
+
     repository: str | None = Field(
         default=None,
         description=(
             "The name of the repository to search in. "
             "If not provided, the search will be performed in all repositories."
         ),
     )
-    intent: str = Field(..., description="A brief description of why you are searching for this code.")
 
 
 class RepositoryStructureInput(BaseModel):

daiv/codebase/search_engines/retrievers.py

@@ -155,7 +155,7 @@ def _get_relevant_documents(self, query: str, *, run_manager: CallbackManagerFor
         )
         return [
             Document(
-                id=document.uuid,
+                id=str(document.uuid),
                 page_content=document.page_content,
                 metadata={
                     "id": str(document.uuid),