fixed llm issues with response

api/routes/chat.py

@@ -36,53 +36,19 @@ async def chat_ai(
         
         # Initialize DroneBot with history and customer metadata
         logger.info("Initializing DroneBot...")
+        logger.info(f"History: {history}")
         bot = DroneBot(history=history, use_openai_as_fallback=True, customer_metadata=customer_metadata)
         logger.info("DroneBot initialized successfully")
         
         # Get response from DroneBot
         logger.info("Calling DroneBot.chat()...")
-        result = bot.chat(request.query)
+        result = await bot.chat(request.query)
         logger.info(f"DroneBot response received: {result}")
         
-        # Validate result and final_message
-        if not result or "final_message" not in result:
-            logger.error(f"Invalid result from DroneBot: {result}")
-            raise HTTPException(
-                status_code=500,
-                detail="Invalid response from DroneBot: missing final_message"
-            )
         
-        final_message = result["final_message"]
-        logger.info(f"Final message extracted: {final_message}")
-        
-        if not final_message or not isinstance(final_message, str):
-            logger.error(f"Final message is not a valid string: {type(final_message)} - {final_message}")
-            raise HTTPException(
-                status_code=500,
-                detail="Invalid response from DroneBot: final_message is not a valid string"
-            )
-        
-        try:
-            logger.info("Attempting to parse final_message as JSON...")
-            message = json.loads(final_message)
-            logger.info("JSON parsing successful")
-        except json.JSONDecodeError as json_error:
-            logger.warning(f"JSON decode error: {json_error}")
-            logger.warning(f"Raw final_message: {final_message}")
-            # If JSON parsing fails, create a fallback structure
-            message = {
-                "message": final_message,
-                "options": None,
-                "requires_selection": False,
-                "end": "in_progress",
-                "form": {}
-            }
-            logger.info("Created fallback message structure")
-        
-        logger.info(f"Final message to return: {message}")
         return ChatResponse(
             status="success",
-            message=message
+            message=result              
         )
     except HTTPException:
         logger.error("Re-raising HTTPException")

src/config/llm_config.py

@@ -3,6 +3,7 @@ class LlmConfig:
         class models:
             gpt_4o = "gpt-4o"
             gpt_4_1 = "gpt-4.1"
+            gpt_5_mini = "gpt-5-mini"
 
         temperatures = {
             "default": 0.7,

src/llm/agent/structured_output.py

@@ -0,0 +1,110 @@
+import os
+import json
+import asyncio
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, SystemMessage
+from typing import List, Optional, Dict, Any
+from typing_extensions import TypedDict, Annotated
+from src.config.llm_config import LlmConfig
+from config import Config
+from logger import logger 
+
+
+class DroneAssessmentResponse(TypedDict):
+    """Structured response for drone assessment that matches the chat template format"""
+    message: Annotated[str, "The main response message about the drone assessment from ai"]
+    options: Annotated[Optional[List[str]], None, "List of options for user selection from ai"]
+    requires_selection: Annotated[bool, False, "Whether the user needs to make a selection from ai"]
+    end: Annotated[str, "in_progress", "The current state: 'in_progress', 'complete', 'cancelled', or other status"]
+    form: Annotated[Dict[str, str], {}, "Form data with assessment results as string key-value pairs"]
+
+
+class structureOuputTool:
+    def __init__(self):
+        self.llm = ChatOpenAI(
+            api_key=Config.OPENAI_API_KEY,
+            model=LlmConfig.openai.models.gpt_4o,
+            temperature=0.3
+        )
+        # Create structured output LLM
+        self.structured_llm = self.llm.with_structured_output(DroneAssessmentResponse)
+
+    def create_assessment_prompt(self, response: dict) -> str:
+        """Create a prompt to convert AI response to structured output"""
+        prompt = f"""
+Your task is to analyze the response from AI and return the structured output in the format provided.
+
+RESPONSE FROM AI:
+{response}
+
+IMPORTANT: Copy the exact values from the response. Do not fabricate or add anything.
+- If message is empty, return empty string
+- If options is empty, return empty list []
+- If form is empty, return empty dict {{}}
+- If requires_selection is empty, return False
+- If end is empty, return "in_progress"
+
+Just convert to structured format - no additional content.
+"""
+        return prompt
+
+    def run(self, booking_form: dict) -> dict:
+        """
+        Convert booking form data to structured assessment response.
+
+        Args:
+            booking_form (dict): Structured booking form input
+
+        Returns:
+            dict: AI-generated structured output
+        """
+        logger.info("Starting DroneAssessmentAgent run...")
+
+        try:
+            # Generate assessment prompt
+            prompt = self.create_assessment_prompt(booking_form)
+            messages = [
+                SystemMessage(content=prompt),
+                HumanMessage(content="Please convert this booking form data into a structured assessment response.")
+            ]
+
+            # Invoke structured LLM
+            logger.debug("Sending prompt to structured LLM...")
+            response = self.structured_llm.invoke(messages)
+
+            # TypedDict response is already a dict, no need to convert
+            result = response
+            logger.info("Received structured LLM response")
+            logger.debug(f"Structured output: {json.dumps(result, indent=2)}")
+
+            return result
+
+        except Exception as e:
+            logger.exception("Error in DroneAssessmentAgent")
+            return {
+                "message": "Error occurred during assessment",
+                "options": ["Try again", "Contact support"],
+                "requires_selection": True,
+                "end": "error",
+                "form": {
+                    "error": str(e),
+                    "assessment_status": "failed"
+                }
+            }
+
+
+async def main():
+    """Async main function to test the drone assessment agent."""
+    from test2 import booking_form_input
+
+    logger.info("Launching DroneAssessmentAgent from main()...")
+    agent = DroneAssessmentAgent()
+    result = await agent.run(booking_form_input)
+
+    logger.info("Drone assessment completed")
+    logger.debug("Final structured output:")
+    logger.debug(json.dumps(result, indent=2))
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

src/llm/orchestrator.py

@@ -16,7 +16,7 @@ from src.prompts.setup_prompt import setup_prompt_manager
 from src.config.llm_config import LlmConfig
 from src.llm.tools import AgenTools
 from config import Config
-
+from src.llm.agent.structured_output import structureOuputTool
 prompt_manager = setup_prompt_manager()
 
 
@@ -43,12 +43,12 @@ class ModelRotationManager:
         # Add OpenAI models if enabled
         if use_openai_as_fallback:
             openai_models = [
-                {
-                    "name": LlmConfig.openai.models.gpt_4o,
-                    "provider": "openai",
-                    "api_key": Config.OPENAI_API_KEY,
-                    "temperature": LlmConfig.openai.temperatures.get("drone_bot")
-                },
+                # {
+                #     "name": LlmConfig.openai.models.gpt_5_mini,
+                #     "provider": "openai",
+                #     "api_key": Config.OPENAI_API_KEY,
+                #     "temperature": LlmConfig.openai.temperatures.get("drone_bot")
+                # },
                 {
                     "name": LlmConfig.openai.models.gpt_4_1,
                     "provider": "openai",
@@ -152,12 +152,14 @@ class DroneBot:
         # Initialize output variables
         self.final_message = ""
         self.final_model_used = ""
+        self.structured_result = None
         
         # Create tools
         self.tools = self._create_tools()
         
         # Initialize prompt manager with customer metadata
         self.prompt_manager = setup_prompt_manager(customer_metadata)
+        self.structure_agent = structureOuputTool()
         
         print(f"DroneBot initialized")
         print(f"OpenAI fallback: {'Enabled' if use_openai_as_fallback else 'Disabled'}")
@@ -206,7 +208,7 @@ class DroneBot:
             
             # Case 3: Response only has tool calls, no text content
             elif hasattr(response, 'tool_calls') and response.tool_calls:
-                return "Generating your visualization..."
+                return "Calling tool..."
             
             # Case 4: Empty or None content
             else:
@@ -265,6 +267,25 @@ class DroneBot:
                     print(f"Extracted final message: {final_message_content}")
                     self.final_message = final_message_content
                     
+                    # Try to parse as JSON first, if it fails, use structured agent
+                    try:
+                        json.loads(final_message_content)
+                        print("Final message is valid JSON, no need for structured agent")
+                    except json.JSONDecodeError:
+                        print("Final message is not valid JSON, calling structured agent")
+                        # Since we can't await here, we'll store the raw message
+                        # The structured processing can happen later if needed
+                        try:
+                            # Try to run synchronously first
+                            structured_result = self.structure_agent.run(self.final_message)
+                            print(f"Structured agent result: {structured_result}")
+                            # Store the structured result separately, keep final_message as string
+                            self.structured_result = structured_result
+                        except Exception as e:
+                            print(f"Error in structured agent (sync): {str(e)}")
+                            # Keep the original message if structured agent fails
+                            pass
+                    
                     # Update state
                     updated_state = {"messages": state["messages"] + [response]}
                     updated_state["current_model"] = current_model_config["name"]
@@ -392,7 +413,7 @@ class DroneBot:
                 langchain_messages.append(AIMessage(content=msg.content))
         return langchain_messages
 
-    def chat(self, user_query: str) -> Dict[str, Any]:
+    async def chat(self, user_query: str) -> Dict[str, Any]:
         """Main method to interact with DroneBot"""
         print(f"DroneBot processing query: {user_query}")
         
@@ -431,28 +452,54 @@ class DroneBot:
                         break
                 
                 if not self.final_message:
-                    self.final_message = "I've processed your visualization request."
+                    self.final_message = "...."
 
-            final_response = {
-                "messages": output.get("messages", []),
-                "final_message": self.final_message,
-                "final_model_used": self.final_model_used or output.get("current_model", "unknown"),
-                "user_question": user_query
-            }
+            # Check if we already have a structured result from the workflow
+            if self.structured_result is not None:
+                print("Using pre-computed structured result from workflow")
+                structured_message = self.structured_result
+            else:
+                # Try to parse final_message as JSON, if it fails, use structured agent
+                try:
+                    structured_message = json.loads(self.final_message)
+                    print("Final message is valid JSON, no structured agent needed")
+                except json.JSONDecodeError:
+                    print("Final message is not valid JSON, calling structured agent")
+                    try:
+                        # Now we can properly await the async function
+                        structured_result = await self.structure_agent.run(self.final_message)
+                        print(f"Structured agent result: {structured_result}")
+                        # Assign the structured result to structured_message
+                        structured_message = structured_result
+                    except Exception as e:
+                        print(f"Error in structured agent: {str(e)}")
+                        structured_message = {
+                            "message": "Sorry i encountered an error while processing your request. Please try again.",
+                            "options": [],
+                            "requires_selection": False,
+                            "end": "error",
+                            "form": {}
+                        }
 
-            print(f"Final message: {self.final_message[:100]}...")
-            return final_response
+            # final_response = {
+            #     "messages": output.get("messages", []),
+            #     "final_message": structured_message,
+            #     "final_model_used": self.final_model_used or output.get("current_model", "unknown"),
+            #     "user_question": user_query
+            # }
+
+            #print(f"Final message: {self.final_message[:100]}...")
+            return structured_message
             
         except Exception as e:
             print(f"Error in DroneBot workflow execution: {str(e)}")
             return {
-                "messages": [],
-                "final_message": "Sorry, I encountered an error while processing your visualization request. Please try again.",
-                "final_model_used": "error",
-                "user_question": user_query,
-                "error": str(e)
-            }
-
+            "message": "Sorry i encountered an error while processing your request. Please try again.",
+            "options": [],
+            "requires_selection": False,
+            "end": "error",
+            "form": {}
+        }
 
 # Example usage
 if __name__ == "__main__":
@@ -477,7 +524,8 @@ if __name__ == "__main__":
     
     query = "Can we start"
     # Chat with DroneBot
-    response = bot.chat(query)
+    import asyncio
+    response = asyncio.run(bot.chat(query))
     
     print("Response:", response["final_message"])
     

src/prompts/templates/chat_templates.py

@@ -709,4 +709,7 @@ NOTE: THIS IS SOLELY YOUR TASK PLEASE, NO OTHER THING
     : DO NOT ADD ```json or ``` or anything else before or after the json format
     : DO NOT ADD ANYTHING ELSE TO THE JSON FORMAT, ONLY THE JSON FORMAT, do not add ```json or ``` or anything else before or after the json format
     CRITICAL !!! ONLY JSON !!!
+
+FINAL CRITICAL INSTRUCTIONS:
+ RETURN ONLY JSON FORMAT NO EXPLANATION OR EXTRA INFORMATION BEFORE OR AFTER, DO NOT ADD ```json before or after
 """

test.py

@@ -1,7 +1,6 @@
-# terminal_chat.py
 from src.llm.orchestrator import DroneBot, Message  # Adjust import path as needed
-
-def terminal_chat():
+import json
+async def terminal_chat():
     print("🚁 DroneBot Terminal Chat")
     print("Type 'exit' to quit.\n")
 
@@ -31,14 +30,16 @@ def terminal_chat():
         bot.history = history
 
         # Get bot response
-        response = bot.chat(user_input)
+        response = await bot.chat(user_input)
 
         # Add bot response to history
-        history.append(Message(role="ai", content=response["final_message"]))
+        history.append(Message(role="ai", content=json.dumps(response)))
 
         # Print bot response
-        print(f"🤖 DroneBot: {response['final_message']}\n")
+        
+        print(f"🤖 DroneBot: {response}\n")
 
 
 if __name__ == "__main__":
-    terminal_chat()
+    import asyncio
+    asyncio.run(terminal_chat())

tests/integration/bot_chat.py

@@ -1,7 +1,7 @@
 # terminal_chat.py
 from src.llm.orchestrator import DroneBot, Message  # Adjust import path as needed
 
-def terminal_chat():
+async def terminal_chat():
     print("🚁 DroneBot Terminal Chat")
     print("Type 'exit' to quit.\n")