fix: Integrate the Agent SDK into the chart generation

Author: Avijit-Microsoft

src/api/agents/chart_agent_factory.py

@@ -0,0 +1,63 @@
+from azure.identity import DefaultAzureCredential
+from azure.ai.projects import AIProjectClient
+
+from agents.agent_factory_base import BaseAgentFactory
+
+
+class ChartAgentFactory(BaseAgentFactory):
+    """
+    Factory class for creating Chart agents that generate chart.js compatible JSON
+    based on numerical and structured data from RAG responses.
+    """
+
+    @classmethod
+    async def create_agent(cls, config):
+        """
+        Asynchronously creates an AI agent configured to convert structured data
+        into chart.js-compatible JSON using Azure AI Project.
+
+        Args:
+            config: Configuration object containing AI project and model settings.
+
+        Returns:
+            dict: A dictionary containing the created 'agent' and its associated 'client'.
+        """
+        instructions = """You are an assistant that helps generate valid chart data to be shown using chart.js with version 4.4.4 compatible.
+        Include chart type and chart options.
+        Pick the best chart type for given data.
+        Do not generate a chart unless the input contains some numbers. Otherwise return a message that Chart cannot be generated.
+        Only return a valid JSON output and nothing else.
+        Verify that the generated JSON can be parsed using json.loads.
+        Do not include tooltip callbacks in JSON.
+        Always make sure that the generated json can be rendered in chart.js.
+        Always remove any extra trailing commas.
+        Verify and refine that JSON should not have any syntax errors like extra closing brackets.
+        Ensure Y-axis labels are fully visible by increasing **ticks.padding**, **ticks.maxWidth**, or enabling word wrapping where necessary.
+        Ensure bars and data points are evenly spaced and not squished or cropped at **100%** resolution by maintaining appropriate **barPercentage** and **categoryPercentage** values."""
+
+        project_client = AIProjectClient(
+            endpoint=config.ai_project_endpoint,
+            credential=DefaultAzureCredential(exclude_interactive_browser_credential=False),
+            api_version=config.ai_project_api_version,
+        )
+
+        agent = project_client.agents.create_agent(
+            model=config.azure_openai_deployment_model,
+            name=f"KM-ChartAgent-{config.solution_name}",
+            instructions=instructions,
+        )
+
+        return {
+            "agent": agent,
+            "client": project_client
+        }
+
+    @classmethod
+    async def _delete_agent_instance(cls, agent_wrapper: dict):
+        """
+        Asynchronously deletes the specified chart agent instance from the Azure AI project.
+
+        Args:
+            agent_wrapper (dict): Dictionary containing the 'agent' and 'client' to be removed.
+        """
+        agent_wrapper["client"].agents.delete_agent(agent_wrapper["agent"].id)

src/api/app.py

@@ -17,6 +17,7 @@
 from agents.conversation_agent_factory import ConversationAgentFactory
 from agents.search_agent_factory import SearchAgentFactory
 from agents.sql_agent_factory import SQLAgentFactory
+from agents.chart_agent_factory import ChartAgentFactory
 from api.api_routes import router as backend_router
 from api.history_routes import router as history_router
 
@@ -34,13 +35,16 @@ async def lifespan(fastapi_app: FastAPI):
     fastapi_app.state.agent = await ConversationAgentFactory.get_agent()
     fastapi_app.state.search_agent = await SearchAgentFactory.get_agent()
     fastapi_app.state.sql_agent = await SQLAgentFactory.get_agent()
+    fastapi_app.state.chart_agent = await ChartAgentFactory.get_agent()
     yield
     await ConversationAgentFactory.delete_agent()
     await SearchAgentFactory.delete_agent()
     await SQLAgentFactory.delete_agent()
+    await ChartAgentFactory.delete_agent()
     fastapi_app.state.sql_agent = None
     fastapi_app.state.search_agent = None
     fastapi_app.state.agent = None
+    fastapi_app.state.chart_agent = None
 
 
 def build_app() -> FastAPI:

src/api/services/chat_service.py

@@ -21,13 +21,13 @@
 from semantic_kernel.agents import AzureAIAgentThread
 from semantic_kernel.exceptions.agent_exceptions import AgentException
 
-from azure.ai.agents.models import TruncationObject
+from azure.ai.agents.models import TruncationObject, MessageRole, ListSortOrder
 
 from cachetools import TTLCache
 
 from helpers.utils import format_stream_response
-from helpers.azure_openai_helper import get_azure_openai_client
 from common.config.config import Config
+from agents.chart_agent_factory import ChartAgentFactory
 
 # Constants
 HOST_NAME = "CKM"
@@ -86,47 +86,59 @@ def __init__(self, request : Request):
         if ChatService.thread_cache is None:
             ChatService.thread_cache = ExpCache(maxsize=1000, ttl=3600.0, agent=self.agent)
 
-    def process_rag_response(self, rag_response, query):
+    async def process_rag_response(self, rag_response, query):
         """
-        Parses the RAG response dynamically to extract chart data for Chart.js.
+        Uses the ChartAgent directly (agentic call) to extract chart data for Chart.js.
         """
         try:
-            client = get_azure_openai_client()
-
-            system_prompt = """You are an assistant that helps generate valid chart data to be shown using chart.js with version 4.4.4 compatible.
-            Include chart type and chart options.
-            Pick the best chart type for given data.
-            Do not generate a chart unless the input contains some numbers. Otherwise return a message that Chart cannot be generated.
-            Only return a valid JSON output and nothing else.
-            Verify that the generated JSON can be parsed using json.loads.
-            Do not include tooltip callbacks in JSON.
-            Always make sure that the generated json can be rendered in chart.js.
-            Always remove any extra trailing commas.
-            Verify and refine that JSON should not have any syntax errors like extra closing brackets.
-            Ensure Y-axis labels are fully visible by increasing **ticks.padding**, **ticks.maxWidth**, or enabling word wrapping where necessary.
-            Ensure bars and data points are evenly spaced and not squished or cropped at **100%** resolution by maintaining appropriate **barPercentage** and **categoryPercentage** values."""
             user_prompt = f"""Generate chart data for -
             {query}
             {rag_response}
             """
-            logger.info(">>> Processing chart data for response: %s", rag_response)
-
-            completion = client.chat.completions.create(
-                model=self.azure_openai_deployment_name,
-                messages=[
-                    {"role": "system", "content": system_prompt},
-                    {"role": "user", "content": user_prompt},
-                ],
-                temperature=0,
+
+            agent_info = await ChartAgentFactory.get_agent()
+            agent = agent_info["agent"]
+            client = agent_info["client"]
+
+            thread = client.agents.threads.create()
+
+            client.agents.messages.create(
+                thread_id=thread.id,
+                role=MessageRole.USER,
+                content=user_prompt
+            )
+
+            run = client.agents.runs.create_and_process(
+                thread_id=thread.id,
+                agent_id=agent.id
             )
 
-            chart_data = completion.choices[0].message.content.strip().replace("```json", "").replace("```", "")
-            logger.info(">>> Generated chart data: %s", chart_data)
+            if run.status == "failed":
+                print(f"[Chart Agent] Run failed: {run.last_error}")
+                return {"error": "Chart could not be generated due to agent failure."}
+
+            chart_json = ""
+            messages = client.agents.messages.list(thread_id=thread.id, order=ListSortOrder.ASCENDING)
+            for msg in messages:
+                if msg.role == MessageRole.AGENT and msg.text_messages:
+                    chart_json = msg.text_messages[-1].text.value.strip()
+                    break
+
+            client.agents.threads.delete(thread_id=thread.id)
+
+            chart_json = chart_json.replace("```json", "").replace("```", "").strip()
+            chart_data = json.loads(chart_json)
+
+            if not chart_data or "error" in chart_data:
+                return {
+                    "error": chart_data.get("error", "Chart could not be generated from this data."),
+                    "hint": "Try asking a question with some numerical values, like 'sales per region' or 'calls per day'."
+                }
 
-            return json.loads(chart_data)
+            return chart_data
 
         except Exception as e:
-            logger.error("Error processing RAG response: %s", e)
+            logger.error("Agent error in chart generation: %s", e)
             return {"error": "Chart could not be generated from this data. Please ask a different question."}
 
     async def stream_openai_text(self, conversation_id: str, query: str) -> StreamingResponse:
@@ -254,7 +266,7 @@ async def complete_chat_request(self, query, last_rag_response=None):
             return {"error": "A previous RAG response is required to generate a chart."}
 
         # Process RAG response to generate chart data
-        chart_data = self.process_rag_response(last_rag_response, query)
+        chart_data = await self.process_rag_response(last_rag_response, query)
 
         if not chart_data or "error" in chart_data:
             return {

src/tests/api/agents/test_chart_agent_factory.py

@@ -0,0 +1,51 @@
+import pytest
+from unittest.mock import AsyncMock, patch, MagicMock
+from agents.chart_agent_factory import ChartAgentFactory
+
+
+@pytest.mark.asyncio
+@patch("agents.chart_agent_factory.DefaultAzureCredential")
+@patch("agents.chart_agent_factory.AIProjectClient")
+async def test_create_agent_success(mock_ai_project_client_class, mock_credential_class):
+    # Mock config
+    mock_config = MagicMock()
+    mock_config.ai_project_endpoint = "https://example-endpoint/"
+    mock_config.ai_project_api_version = "2024-04-01-preview"
+    mock_config.azure_openai_deployment_model = "gpt-4"
+    mock_config.solution_name = "TestSolution"
+
+    # Mock client and agent
+    mock_agent = MagicMock()
+    mock_client = MagicMock()
+    mock_client.agents.create_agent.return_value = mock_agent
+    mock_ai_project_client_class.return_value = mock_client
+    mock_credential_class.return_value = MagicMock()
+
+    # Call create_agent
+    result = await ChartAgentFactory.create_agent(mock_config)
+
+    # Assertions
+    assert result["agent"] == mock_agent
+    assert result["client"] == mock_client
+    mock_ai_project_client_class.assert_called_once_with(
+        endpoint=mock_config.ai_project_endpoint,
+        credential=mock_credential_class.return_value,
+        api_version=mock_config.ai_project_api_version
+    )
+    mock_client.agents.create_agent.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_delete_agent_instance():
+    mock_client = MagicMock()
+    mock_agent = MagicMock()
+    mock_agent.id = "mock-agent-id"
+    
+    agent_wrapper = {
+        "agent": mock_agent,
+        "client": mock_client
+    }
+
+    await ChartAgentFactory._delete_agent_instance(agent_wrapper)
+
+    mock_client.agents.delete_agent.assert_called_once_with("mock-agent-id")