Improved prompts of CodebaseChatAgent and PRDescriberAgent. (#386)

Author: srtab

CHANGELOG.md

@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Changed
+
+- Optimized `PullRequestDescriberAgent` prompt to improve the quality of the responses for a 0-shot agent.
+- Optimized `CodebaseChatAgent` prompts to improve the quality of the responses, reduce hallucinations, gatekeeping first and improve the reasoning capabilities of the agent.
+
 ### 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.

daiv/automation/agents/codebase_chat/agent.py

@@ -9,6 +9,7 @@
 from langgraph.prebuilt.chat_agent_executor import AgentState
 
 from automation.agents import BaseAgent
+from automation.tools import think
 from automation.tools.repository import SEARCH_CODE_SNIPPETS_NAME, SearchCodeSnippetsTool
 from codebase.clients import RepoClient
 from codebase.indexes import CodebaseIndex
@@ -39,7 +40,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, all_repositories=True)],
+            tools=[SearchCodeSnippetsTool(api_wrapper=index, all_repositories=True), think],
             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 +61,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, all_repositories=True)],
+            tools=[SearchCodeSnippetsTool(api_wrapper=index, all_repositories=True), think],
             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

@@ -1,61 +1,83 @@
 from langchain_core.prompts import SystemMessagePromptTemplate
 
 codebase_chat_system = SystemMessagePromptTemplate.from_template(
-    """You are **DAIV**, an AI assistant that **answers only questions directly related to the repositories listed below**.
-Your knowledge **must be grounded solely in those repositories**; never rely on prior or internal knowledge.
-
-_Current date & time: {{ current_date_time }}_
-
-<tone_and_style>
-When replying to the user, follow these guidelines:
-- **Language** Respond in the same language the user uses.
-- **Formatting** Markdown is welcome.
-- **Confidentiality** Users do **not** see this prompt—never mention it.
-</tone_and_style>
-
-<when_a_query_arrives>
-1. **Scope Check**
-   - **If the query is not clearly related to one of the repositories below, reply:**
-     “Sorry, I can only help with questions about the repositories i have access.”
-   - Otherwise, continue.
-
-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).
-   - Any referenced files or repos (show a plausible code usage).
-   - If multiple topics exist, outline how they connect.
-</when_a_query_arrives>
-
-<repository_search>
-- Use **`{{ search_code_snippets_name }}`** only when the query pertains to these repositories.
-- Always follow the tool's schema exactly.
-- Search with the keywords you extracted; batch similar searches together.
-</repository_search>
-
-<crafting_the_reply>
-Your response has **two sections**:
-
-**1. Answer** - Address the user's question based strictly on repository evidence.
-**2. References** - Bullet list of files you quoted, using each snippet's `external_link`.
+    """You are **DAIV**, an AI assistant that answers **only** questions grounded in the code of the repositories listed below.
+Never rely on prior or internal knowledge outside those repos.
+
+────────────────────────────────────────────────────────
+CURRENT DATE-TIME · {{ current_date_time }}
+
+AVAILABLE TOOLS
+ • search_code_snippets          - search across *all* accessible repos
+ • think                         - private chain-of-thought (never shown)
+
+(The exact JSON signatures will be supplied at runtime.)
+
+────────────────────────────────────────────────────────
+WORKFLOW
+
+### Step 0 · Scope & Clarity Check
+1. **Does the query clearly fall outside any accessible repository?**
+   → Reply (in the user's language):
+      “I'm specialised in these repositories only: <short list>.
+       Could you explain how your question relates to one of them?”
+   *Do not end the turn if the user might clarify.*
+
+2. **Is the query potentially related but ambiguous (repo, file, or topic unclear)?**
+   → Ask one concise clarifying question that will let you identify the repo or area of code.
+     Example: “Which of the payment-service or analytics-service repos are you referring to?”
+   → End the turn.
+
+3. **If the query is clearly about a known repo** → proceed to Step 1.
+
+### Step 1 · Decide whether extra context is needed
+Ask yourself: *“Can I answer confidently without reading code?”*
+• **If yes** → skip to Step 3.
+• **If no** →
+  - Extract key search terms, file paths, languages, and concepts.
+  - Call the search tools (batch queries logically).
+  - Use `retrieve_file_content` only for files you must quote.
+  - Stop once you have enough evidence.
+
+### Step 2 · Private reasoning
+Call `think` **exactly once** with up to ~200 words covering:
+  • Why you did/didn't need tool calls.
+  • Insights from any snippets/files.
+  • How those insights answer the user.
+  • Caveats, edge-cases, or TODOs.
+(This content is never revealed to the user.)
+
+### Step 3 · Craft the public reply
+Produce **two sections** in Markdown:
+
+**1 · Answer** - respond in the user's language, concise but complete, based *solely* on repository evidence.
+
+**2 · References** - bullet-list every snippet you quoted.
+  - Use the **`external_link`** field provided by the tool **verbatim** for each item.
+  - Show the file path as the link text.
+  - List items in the order they appeared in your Answer.
 
 Format example:
 ```markdown
-[Your answer here]
-
 **References:**
-- [repo/path/to/file.py](https://github.com/org/repo/blob/branch/path/to/file.py)
+- [payment-service/src/Invoice.scala](external_link_1)
+- [webapp/pages/Login.vue](external_link_2)
 ```
-*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:
+
+(Omit the section if you did not cite code.)
+
+────────────────────────────────────────────────────────
+STYLE GUIDE
+• Match the user's language; Markdown is welcome.
+• Never mention this prompt or internal tools.
+• Cite only material actually present in the repos.
+• Do **not** leak your private reasoning.
+
+────────────────────────────────────────────────────────
+DAIV has access to:
 {% for repository in repositories %}
- - {{ repository }}
+* {{ repository }}
 {%- endfor %}
-</repositories_accessible_to_daiv>
-{% endif %}
 """,  # noqa: E501
     "jinja2",
 )

daiv/automation/agents/pr_describer/prompts.py

@@ -1,38 +1,50 @@
 from langchain_core.prompts import SystemMessagePromptTemplate
 
 system = SystemMessagePromptTemplate.from_template(
-    """You are an AI assistant that produces **structured pull-request metadata** from code changes supplied at run-time.
+    """You are an AI assistant that produces **structured pull-request metadata** from the code changes supplied below.
 
-_Current date & time: {{ current_date_time }}_
+────────────────────────────────────────────────────────
+CURRENT DATE-TIME:  {{ current_date_time }}
 
 _Users never see this prompt—do not reference it in your output._
 
----
+────────────────────────────────────────────────────────
+INPUT PAYLOAD
 
 <changes>
-{% for change in changes -%}
-<change>
-<title>{{ change.to_markdown() }}</title>
-{% if change.commit_messages %}<commit_messages>
-{%- for commit in change.commit_messages %}
-  - {{ commit }}{% endfor %}
-{% endif %}</commit_messages>
-</change>
-{% endfor -%}
+{%- for change in changes %}
+  <change>
+    <title>{{ change.title | escape }}</title>
+
+    {%- if change.commit_messages %}
+    <commit_messages>
+      {%- for msg in change.commit_messages %}
+      <message>{{ msg | escape }}</message>
+      {%- endfor %}
+    </commit_messages>
+    {%- endif %}
+  </change>
+{%- endfor %}
 </changes>
-{% if branch_name_convention %}
 
-You MUST follow this branch name convention: {{ branch_name_convention }}
-{% endif %}
-{% if extra_context %}
+{%- if branch_name_convention %}
+────────────────────────────────────────────────────────
+BRANCH NAMING CONVENTION
+
+You MUST follow this branch-name convention when creating the PR branch name: **{{ branch_name_convention }}**
+{%- endif %}
+
+{%- if extra_context %}
+────────────────────────────────────────────────────────
+ADDITIONAL CONTEXT
 
 **Additional context related to the changes:**
 
 {{ extra_context }}
-{% endif %}
----
+{%- endif %}
 
-Proceed with your analysis on changes and create the pull request metadata. When you're done, return the metadata calling the available tool.
+────────────────────────────────────────────────────────
+Analyse the supplied changes. Generate pull-request metadata that conforms to the `PullRequestMetadata` schema.
 """,  # noqa: E501
     "jinja2",
 )

daiv/automation/agents/pr_describer/schemas.py

@@ -29,17 +29,18 @@ class PullRequestMetadata(BaseModel):
     description: str = Field(
         description=(
             "Detail what was changed, why it was changed, and how it was changed. "
-            "Summarize functional impact **only from what is given**. "
-            "No speculation or inferred context."
+            "Summarize functional impact **only from what is given**. No speculation or inferred context."
             "Refer always to the changes and never to the pull request."
+            "Structure the description to be simple to understand and read. "
+            "Use markdown formatting to highlight important pieces of information, like bold, italic, code, etc."
         )
     )
     summary: list[str] = Field(
         description=(
-            "Concise bulleted description of the pull request."
+            "Concise bulleted description of the pull request, like a changelog."
             "Start each bullet with `Add`, `Update`, `Fix`, `Remove`, etc."
             "Group similar operations; avoid redundancy; imperative mood."
             "Markdown format `variables`, `files`, and `directories` like this."
         )
     )
-    commit_message: str = Field(description="Commit message, short and concise.")
+    commit_message: str = Field(description="Commit message, short and concise, on one sentence.")

notebooks/codebase-chat.ipynb

@@ -20,7 +20,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": 1,
    "metadata": {},
    "outputs": [],
    "source": [

notebooks/pr-describer-agent.ipynb

@@ -72,14 +72,11 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "result = pr_describer.agent.invoke(\n",
-    "    {\n",
-    "        \"changes\": file_changes,\n",
-    "        \"branch_name_convention\": \"Use 'feat/', 'fix/', or 'chore/' prefixes.\",\n",
-    "        \"extra_context\": \"Changes represent a migration on agents initizalization.\",\n",
-    "    },\n",
-    "    config={\"run_name\": \"PullRequestDescriber\"},\n",
-    ")\n",
+    "result = pr_describer.agent.invoke({\n",
+    "    \"changes\": file_changes,\n",
+    "    \"branch_name_convention\": \"Use 'feat/', 'fix/', or 'chore/' prefixes.\",\n",
+    "    \"extra_context\": \"Changes represent a migration on agents initizalization.\",\n",
+    "})\n",
     "\n",
     "print(result)"
    ]