Add agent workflow with multi-agent routing and query processing

Author: shakil1819

app/v1/routers/agent_workflow/__init__.py

@@ -0,0 +1,173 @@
+from typing import Any, Literal, Optional  # noqa: UP035
+
+from pydantic import BaseModel
+from pydantic_graph import BaseNode, End, Graph, GraphRunContext
+
+from app.tools.data_analyst_agent import DataAnalystAgent
+from app.tools.sql_data_analyst_agent import SqlDataAnalystAgent
+from app.v1.routers.agent_workflow.persistence import ChromaDBStatePersistence, get_chroma_client
+
+
+class AgentState(BaseModel):
+    """State model for the agent workflow graph."""
+
+    query: str
+    result: Optional[dict[str, Any]] = None
+    agent_used: Optional[str] = None
+
+
+class QueryClassifier(BaseNode[AgentState]):
+    """Node to classify the query and determine which agent to use."""
+
+    async def run(self, ctx: GraphRunContext[AgentState]) -> Literal["use_sql_agent", "use_data_agent"]:
+        """Classify the query to determine which agent to use.
+
+        Returns:
+            str: Either "use_sql_agent" or "use_data_agent"
+
+        """
+        # Simple classifier based on keywords
+        query = ctx.state.query.lower()
+
+        # Keywords that suggest SQL operations
+        sql_keywords = [
+            "sql",
+            "query",
+            "table",
+            "database",
+            "select",
+            "join",
+            "where",
+            "group by",
+            "order by",
+            "having",
+            "count",
+            "sum",
+            "average",
+            "filter",
+            "find records",
+            "search records",
+        ]
+
+        # Keywords that suggest data analysis or visualization
+        data_keywords = [
+            "analyze",
+            "chart",
+            "graph",
+            "plot",
+            "visualization",
+            "trend",
+            "pattern",
+            "correlation",
+            "distribution",
+            "histogram",
+            "scatter plot",
+            "bar chart",
+            "pie chart",
+            "dashboard",
+        ]
+
+        # Count matches for each category
+        sql_score = sum(1 for keyword in sql_keywords if keyword in query)
+        data_score = sum(1 for keyword in data_keywords if keyword in query)
+
+        # Determine which agent to use based on keyword matches
+        if sql_score > data_score:
+            return "use_sql_agent"
+        return "use_data_agent"
+
+
+class UseSqlAgent(BaseNode[AgentState]):
+    """Node to process the query using the SQL Data Analyst Agent."""
+
+    async def run(self, ctx: GraphRunContext[AgentState]) -> End[dict[str, Any]]:
+        """Process the query using the SQL Data Analyst Agent.
+
+        Returns:
+            End: The result from the SQL Data Analyst Agent
+
+        """
+        sql_agent = SqlDataAnalystAgent()
+        result = await sql_agent.process_query(ctx.state.query)
+
+        # Update the state
+        ctx.state.result = result
+        ctx.state.agent_used = "sql_data_analyst"
+
+        # Return the result and end the graph execution
+        return End({"result": result, "agent_used": "sql_data_analyst"})
+
+
+class UseDataAgent(BaseNode[AgentState]):
+    """Node to process the query using the Data Analyst Agent."""
+
+    async def run(self, ctx: GraphRunContext[AgentState]) -> End[dict[str, Any]]:
+        """Process the query using the Data Analyst Agent.
+
+        Returns:
+            End: The result from the Data Analyst Agent
+
+        """
+        data_agent = DataAnalystAgent()
+        result = data_agent.process_query(ctx.state.query)
+
+        # Update the state
+        ctx.state.result = result
+        ctx.state.agent_used = "data_analyst"
+
+        # Return the result and end the graph execution
+        return End({"result": result, "agent_used": "data_analyst"})
+
+
+class AgentWorkflowGraph:
+    """Multi-agent workflow graph implementation using pydantic-ai and ChromaDB.
+
+    This class manages the workflow that selects and runs the appropriate agent
+    based on the user's query.
+    """
+
+    def __init__(self):
+        """Initialize the agent workflow graph."""
+        # Create the graph with nodes and edges
+        nodes = {
+            "start": QueryClassifier(),
+            "use_sql_agent": UseSqlAgent(),
+            "use_data_agent": UseDataAgent(),
+        }
+
+        edges = {
+            "start": {
+                "use_sql_agent": "use_sql_agent",
+                "use_data_agent": "use_data_agent",
+            }
+        }
+
+        # Create the graph using named parameters
+        self.graph = Graph(
+            state_type=AgentState,
+            nodes=nodes,
+            edges=edges,
+        )
+
+        # Set up persistence with ChromaDB
+        chroma_client = get_chroma_client()
+        self.persistence = ChromaDBStatePersistence(
+            chroma_client=chroma_client,
+            collection_name="agent_workflow_states",
+        )
+
+    async def process_query(self, query: str) -> dict[str, Any]:
+        """Process a user query through the multi-agent workflow.
+
+        Args:
+            query: The user's query string
+
+        Returns:
+            dict: A dictionary containing the result and the agent used
+
+        """
+        # Initialize the state with the query
+        state = AgentState(query=query)
+
+        # Run the graph with the state
+        return await self.graph.arun(state, state_persistence=self.persistence)

app/v1/routers/agent_workflow/graph.py

@@ -0,0 +1,105 @@
+import json
+import os
+import uuid
+
+from typing import Any, Optional  # noqa: UP035
+
+import chromadb
+
+from chromadb.config import Settings
+from pydantic_graph.persistence import BaseStatePersistence
+
+
+def get_chroma_client(persist_directory: Optional[str] = None) -> chromadb.Client:
+    """Get or create a ChromaDB client.
+
+    Args:
+        persist_directory: Optional directory to persist the ChromaDB data.
+            If not provided, an in-memory client will be used.
+
+    Returns:
+        chromadb.Client: A ChromaDB client instance
+
+    """
+    if persist_directory:
+        # Create the directory if it doesn't exist
+        os.makedirs(persist_directory, exist_ok=True)
+
+        # Return a persistent client
+        return chromadb.PersistentClient(path=persist_directory, settings=Settings(anonymized_telemetry=False))
+    # Return an in-memory client
+    return chromadb.Client(Settings(anonymized_telemetry=False, is_persistent=False))
+
+
+class ChromaDBStatePersistence(BaseStatePersistence):
+    """Custom implementation of state persistence using ChromaDB."""
+
+    def __init__(
+        self,
+        chroma_client: chromadb.Client,
+        collection_name: str = "graph_states",
+    ):
+        """Initialize the ChromaDB state persistence.
+
+        Args:
+            chroma_client: ChromaDB client instance
+            collection_name: Name of the collection to store states
+
+        """
+        self.client = chroma_client
+        self.collection = self.client.get_or_create_collection(collection_name)
+
+    async def save_state(self, state_id: str, state: dict[str, Any], metadata: Optional[dict[str, Any]] = None) -> None:
+        """Save a state to ChromaDB.
+
+        Args:
+            state_id: Unique identifier for the state
+            state: The state to save
+            metadata: Optional metadata to store with the state
+
+        """
+        # Convert state to JSON string
+        state_json = json.dumps(state)
+
+        # Prepare metadata
+        meta = metadata or {}
+        meta["state_id"] = state_id
+
+        # Use a unique document ID
+        doc_id = f"{state_id}_{uuid.uuid4()}"
+
+        # Add the state to ChromaDB
+        self.collection.add(documents=[state_json], metadatas=[meta], ids=[doc_id])
+
+    async def load_state(self, state_id: str) -> Optional[dict[str, Any]]:
+        """Load a state from ChromaDB.
+
+        Args:
+            state_id: Unique identifier for the state
+
+        Returns:
+            The state if found, None otherwise
+
+        """
+        # Query ChromaDB for the state
+        results = self.collection.query(where={"state_id": state_id}, limit=1)
+
+        # Return the state if found
+        if results["documents"] and results["documents"][0]:
+            return json.loads(results["documents"][0])
+
+        return None
+
+    async def delete_state(self, state_id: str) -> None:
+        """Delete a state from ChromaDB.
+
+        Args:
+            state_id: Unique identifier for the state
+
+        """
+        # Find all documents with the given state_id
+        results = self.collection.query(where={"state_id": state_id}, limit=100)
+
+        # Delete the documents if found
+        if results["ids"]:
+            self.collection.delete(ids=results["ids"])

app/v1/routers/agent_workflow/persistence.py

@@ -0,0 +1,33 @@
+from fastapi import APIRouter, HTTPException
+from pydantic import BaseModel
+
+from app.v1.routers.agent_workflow.graph import AgentWorkflowGraph
+
+
+router = APIRouter()
+agent_workflow = AgentWorkflowGraph()
+
+
+class QueryRequest(BaseModel):
+    """Request model for agent workflow queries."""
+
+    query: str
+
+
+class QueryResponse(BaseModel):
+    """Response model for agent workflow results."""
+
+    result: dict
+    agent_used: str
+
+
+@router.post("/query", response_model=QueryResponse)
+async def process_query(request: QueryRequest) -> QueryResponse:
+    """Process a user query through the multi-agent workflow.
+
+    The workflow will automatically select the appropriate agent based on the query.
+    """
+    try:
+        return await agent_workflow.process_query(request.query)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Error processing query: {str(e)}") from e

app/v1/routers/agent_workflow/router.py

@@ -1,7 +1,9 @@
 from fastapi import APIRouter
 
+from app.v1.routers.agent_workflow import router as agent_workflow_router
 from app.v1.routers.users import users_router
 
 
 router = APIRouter(prefix="/v1")
 router.include_router(users_router.router, prefix="/users", tags=["Users"])
+router.include_router(agent_workflow_router.router, prefix="/agent-workflow", tags=["Agent Workflow"])

app/v1/routers/base_router.py

@@ -9,6 +9,7 @@ requires-python = ">=3.9"
 dependencies = [
   "agno>=1.2.6",
   "ai-data-science-team>=0.0.0.9016",
+  "chromadb>=1.0.0",
   "commitizen>=4.4.1",
   "duckdb>=1.2.1",
   "fastapi[standard]>=0.115.8",
@@ -27,6 +28,7 @@ dependencies = [
   "phidata>=2.7.10",
   "pre-commit>=4.1.0",
   "pydantic-ai[logfire]>=0.0.46",
+  "pydantic-graph>=0.0.52",
   "scikit-learn>=1.6.1",
   "streamlit>=1.44.1",
 ]