Enhance question-answering API with agent types and streaming support

jenilsavani9 · jenilsavani9 · commit f55f2cfa9732 · 2025-12-04T23:54:06.000+05:30
diff --git a/apps/api/routes/v1/ask.py b/apps/api/routes/v1/ask.py
@@ -5,19 +5,13 @@
 from fastapi import APIRouter, HTTPException, status
 from fastapi.responses import StreamingResponse
 
-from apps.api.routes.v1.schemas.qa import (
-    QARequest,
-    ContextItem,
-    IssueRequest,
-    IssueContextItem,
-)
+from apps.api.routes.v1.schemas.qa import QARequest, ContextItem
 from apps.api.utils import APIResponse
 from packages.memory.qa_service import QAService
 from packages.memory.services.graph_service import GraphService
-from packages.memory.services.issue_analysis import IssueAnalyzer
 from packages.database.graph.graph import neo4j_client
 from packages.config import Settings
-from packages.config.feature_flags import is_feature_enabled, FeatureFlag
+
 
 router = APIRouter(prefix="/ask", tags=["Question Answering"])
 settings = Settings()
@@ -26,44 +20,48 @@
 # Initialize services once per module import
 graph_service = GraphService(neo4j_client)
 qa_service = QAService(graph_service)
-issue_analyzer = IssueAnalyzer(graph_service)
 
 
 @router.post(
     "",
     summary="Ask a question about the codebase",
     description=(
-        "Submit a natural language question about the codebase. The system "
-        "retrieves relevant code context (hybrid or vector search) and uses an LLM "
-        "to generate an answer."
+        "Submit a question about the codebase with different agent types. "
+        "Supports streaming responses and specialized agents: "
+        "pathfinder (code structure), chronicle (history), diagnostician (debugging), "
+        "blueprint (architecture), sentinel (code review)."
     ),
 )
-async def question_answer(request: QARequest):
-    """Return an AI-generated answer with supporting context.
-
-    Process:
-    1. Retrieve relevant code elements using the selected search strategy.
-    2. Provide context snippets to the configured LLM provider.
-    3. Return an answer plus the context used.
-    """
+async def ask(request: QARequest):
+    """Ask a question using different agent types."""
     try:
-        logger.info(f"Processing question: {request.question[:50]}...")
+        logger.info(
+            f"Processing question with {request.agent_type.value} agent: {request.question[:50]}..."
+        )
 
+        # If streaming is requested, return streaming response
+        if request.stream:
+            return StreamingResponse(
+                _stream_answer(request),
+                media_type="text/event-stream"
+            )
+
+        # Non-streaming response
         result: dict[str, Any] = qa_service.ask(
             question=request.question,
+            agent_type=request.agent_type.value,
             search_type=request.search_type,
-            node_type=request.node_type,
             context_limit=request.context_limit,
         )
 
         return APIResponse.success(
             data={
                 "answer": result["answer"],
                 "question": result["question"],
+                "agent_type": result["agent_type"],
                 "context": [ContextItem(**item) for item in result["context"]],
                 "context_count": result["context_count"],
                 "search_type": result["search_type"],
-                "node_type": result.get("node_type"),
                 "node_types": result.get("node_types"),
                 "model": result["model"],
                 "provider": result["provider"],
@@ -84,72 +82,25 @@ async def question_answer(request: QARequest):
             detail=f"Service error: {str(e)}",
         )
     except Exception as e:  # noqa: BLE001
-        logger.exception("Unexpected error in question_answer endpoint")
+        logger.exception("Unexpected error in ask endpoint")
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail="An unexpected error occurred while processing your question",
         )
 
 
-@router.post(
-    "/issue",
-    summary="Analyze a GitHub issue",
-    description=(
-        "Analyze a GitHub issue to identify causes and solutions based on the Knowledge Graph. "
-        "Returns a detailed report and the context used."
-    ),
-)
-async def analyze_issue(request: IssueRequest):
-    """Analyze an issue and return a report.
-
-    Process:
-    1. Search for relevant code and commit history using hybrid search.
-    2. Use LLM to generate a root cause analysis and suggested fix.
-    3. Return the report and context.
-    """
-    try:
-        logger.info(f"Analyzing issue: {request.title}")
-
-        if is_feature_enabled(FeatureFlag.ENABLE_TOKEN_STREAMING):
-            return StreamingResponse(
-                _stream_issue_analysis(request.title, request.body),
-                media_type="text/event-stream"
-            )
-
-        result = issue_analyzer.analyze_issue(request.title, request.body)
-
-        if "error" in result:
-            raise HTTPException(
-                status_code=status.HTTP_404_NOT_FOUND,
-                detail=result["error"],
-            )
-
-        return APIResponse.success(
-            data={
-                "report": result["report"],
-                "context_used": [IssueContextItem(**item) for item in result["context_used"]],
-            },
-            message="Successfully analyzed the issue."
-        )
-
-    except HTTPException:
-        raise
-    except Exception as e:
-        logger.exception("Unexpected error in analyze_issue endpoint")
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail="An unexpected error occurred. Please try again later.",
-        )
-
-
-def _stream_issue_analysis(title: str, body: str):
-    """Generator for streaming issue analysis."""
+def _stream_answer(request: QARequest):
+    """Generator for streaming answers."""
     try:
-        for chunk in issue_analyzer.analyze_issue_stream(title, body):
+        for chunk in qa_service.ask_stream(
+            question=request.question,
+            agent_type=request.agent_type.value,
+            search_type=request.search_type,
+            context_limit=request.context_limit,
+        ):
             yield f"data: {json.dumps(chunk)}\n\n"
     except Exception as e:
-        logger.exception("Error during streaming issue analysis")
+        logger.exception("Error during streaming")
         yield f"data: {json.dumps({'error': str(e)})}\n\n"
     finally:
         yield "data: [DONE]\n\n"
- 
diff --git a/apps/api/routes/v1/schemas/qa.py b/apps/api/routes/v1/schemas/qa.py
@@ -5,35 +5,48 @@
 from pydantic import BaseModel, Field
 from typing import List, Optional, Dict, Any
 
+from packages.memory.agents import AgentType
+
 
 class QARequest(BaseModel):
     """Request model for question-answering."""
     
     question: str = Field(
         ...,
         min_length=1,
-        max_length=500,
         description="The question to answer",
         examples=["How does the authentication system work?"]
     )
     
+    agent_type: AgentType = Field(
+        default=AgentType.PATHFINDER,
+        description=(
+            "Type of agent to use: "
+            "pathfinder (code structure), "
+            "chronicle (commit history), "
+            "diagnostician (debugging), "
+            "blueprint (architecture), "
+            "sentinel (code review)"
+        )
+    )
+    
     search_type: str = Field(
         default="hybrid",
         pattern="^(vector|hybrid)$",
         description="Type of search to use for context retrieval"
     )
     
-    node_type: str = Field(
-        default="Commit",  # Changed from "Function" to "Commit" to match available data
-        description="Type of code nodes to search (Function, Class, File, Commit, etc.)"
-    )
-    
     context_limit: int = Field(
         default=5,
         ge=1,
         le=20,
         description="Maximum number of context items to retrieve"
     )
+    
+    stream: bool = Field(
+        default=False,
+        description="Whether to stream the response"
+    )
 
 
 
@@ -52,39 +65,10 @@ class QAResponse(BaseModel):
     
     answer: str = Field(description="Natural language answer from LLM")
     question: str = Field(description="The original question")
+    agent_type: AgentType = Field(description="Type of agent used")
     context: List[ContextItem] = Field(description="Context items used for the answer")
     context_count: int = Field(description="Number of context items retrieved")
     search_type: str = Field(description="Type of search used")
-    node_type: Optional[str] = Field(None, description="Node type searched")
-    node_types: Optional[List[str]] = Field(None, description="Node types searched (multi-type)")
+    node_types: List[str] = Field(description="Node types searched")
     model: str = Field(description="LLM model used")
     provider: str = Field(description="LLM provider used")
-
-
-class IssueRequest(BaseModel):
-    """Request model for issue analysis."""
-    title: str = Field(
-        ...,
-        min_length=1,
-        max_length=500,
-        description="Title of the issue"
-    )
-    body: str = Field(
-        ...,
-        min_length=1,
-        max_length=500,
-        description="Body/Description of the issue"
-    )
-class IssueContextItem(BaseModel):
-    """Context item for issue analysis."""
-    type: str
-    name: str
-    score: Optional[float] = None
-
-
-class IssueResponse(BaseModel):
-    """Response model for issue analysis."""
-    report: str = Field(description="Analysis report generated by LLM")
-    context_used: List[IssueContextItem] = Field(description="Context items used for analysis")
-
-
diff --git a/packages/memory/agents.py b/packages/memory/agents.py
@@ -0,0 +1,40 @@
+from enum import Enum
+from typing import List
+
+
+class AgentType(str, Enum):    
+    PATHFINDER = "pathfinder"       # Code structure & navigation
+    CHRONICLE = "chronicle"         # Commit history & evolution
+    DIAGNOSTICIAN = "diagnostician" # Debugging & error analysis
+    BLUEPRINT = "blueprint"         # Architecture reasoning
+    SENTINEL = "sentinel"           # Code review & quality checks
+    
+    @classmethod
+    def list_values(cls) -> list[str]:
+        return [agent.value for agent in cls]
+    
+    @classmethod
+    def get_description(cls, agent_type: "AgentType") -> str:
+        descriptions = {
+            cls.PATHFINDER: "Code structure & navigation",
+            cls.CHRONICLE: "Commit history & evolution",
+            cls.DIAGNOSTICIAN: "Debugging & error analysis",
+            cls.BLUEPRINT: "Architecture reasoning",
+            cls.SENTINEL: "Code review & quality checks",
+        }
+        return descriptions.get(agent_type, "Unknown agent type")
+
+    @classmethod
+    def get_node_types(cls, agent_type: str) -> List[str]:
+        """
+        Get the relevant node types to search for each agent.
+        Returns a list of node types ordered by priority.
+        """
+        node_types_map = {
+            cls.PATHFINDER.value: ["Function", "Class", "File"],
+            cls.CHRONICLE.value: ["Commit", "File"],
+            cls.DIAGNOSTICIAN.value: ["Function", "Class", "Commit", "File"],
+            cls.BLUEPRINT.value: ["Class", "File", "Function"],
+            cls.SENTINEL.value: ["Function", "Class", "File", "Commit"],
+        }
+        return node_types_map.get(agent_type, ["Function", "Class", "File"])
diff --git a/packages/memory/prompts.py b/packages/memory/prompts.py
@@ -0,0 +1,78 @@
+from typing import Dict
+from packages.memory.agents import AgentType
+
+class PromptFactory:    
+    PROMPTS: Dict[str, str] = {
+        AgentType.PATHFINDER.value: """
+          You are Pathfinder, an expert code navigator and structure analyst.
+          You help developers understand codebase structure, find files, classes, functions, and navigate complex codebases.
+
+          Your task:
+          1. Help users find specific code elements (functions, classes, files)
+          2. Explain code organization and module structure
+          3. Trace code paths and dependencies
+          4. Guide navigation through the codebase
+          5. Reference specific file paths and function names
+
+          Be precise with locations. Always mention file paths and line references when possible.""",
+        
+        AgentType.CHRONICLE.value: """
+          You are Chronicle, an expert in code history and evolution.
+          You analyze commit history, track changes over time, and help understand how code evolved.
+
+          Your task:
+          1. Explain what changed and when
+          2. Identify who made specific changes
+          3. Track the evolution of features or bugs
+          4. Connect commits to code changes
+          5. Provide timeline context for code decisions
+
+          Reference commit hashes, dates, and authors when available.""",
+        
+        AgentType.DIAGNOSTICIAN.value: """
+          You are Diagnostician, an expert debugger and error analyst.
+          You help identify bugs, analyze errors, and find root causes of issues.
+
+          Your task:
+          1. **Root Cause Analysis**: What is likely causing the issue?
+          2. **Affected Areas**: Which files, classes, or functions are involved?
+          3. **Suggested Fix**: How can this be fixed? Provide code snippets if possible.
+          4. **Relevant History**: Are there recent commits that might have introduced this?
+
+          Be specific. Reference filenames, function names, and line numbers.
+          Think step-by-step through the error flow.""",
+        
+        AgentType.BLUEPRINT.value: """
+          You are Blueprint, an expert software architect and design analyst.
+          You help understand system architecture, design patterns, and high-level code organization.
+
+          Your task:
+          1. Explain architectural decisions and patterns
+          2. Describe how components interact
+          3. Identify design patterns in use
+          4. Suggest architectural improvements
+          5. Map out system dependencies and data flow
+
+          Think at the system level. Explain the "why" behind design choices.""",
+        
+        AgentType.SENTINEL.value: """
+          You are Sentinel, an expert code reviewer and quality analyst.
+          You help identify code quality issues, suggest improvements, and enforce best practices.
+
+          Your task:
+          1. Identify potential bugs or code smells
+          2. Suggest improvements for readability and maintainability
+          3. Check for security vulnerabilities
+          4. Recommend testing strategies
+          5. Enforce coding standards and best practices
+
+          Be constructive and specific. Provide actionable feedback with examples."""
+    }
+    
+    @classmethod
+    def get_prompt(cls, agent_type: str) -> str:
+        return cls.PROMPTS.get(agent_type.lower(), cls.PROMPTS[AgentType.PATHFINDER.value])
+    
+    @classmethod
+    def get_available_types(cls) -> list[str]:
+        return AgentType.list_values()
diff --git a/packages/memory/qa_service.py b/packages/memory/qa_service.py
