fixed llm issues with response

api/routes/chat.py

@@ -42,24 +42,13 @@ async def chat_ai(
         
         # 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}")
         
-        final_message_json = json.loads(result["final_message"])
-            # If JSON parsing fails, create a fallback structure
-        # message = {
-        #         "message": final_message_json,
-        #         "options": None,
-        #         "requires_selection": False,
-        #         "end": "in_progress",
-        #         "form": {}
-        #     }
-        logger.info("Created fallback message structure")
         
-        logger.info(f"Final message to return: {final_message_json}")
         return ChatResponse(
             status="success",
-            message=final_message_json              
+            message=result              
         )
     except HTTPException:
         logger.error("Re-raising HTTPException")

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()
 
 
@@ -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 = {
+            # Check if we already have a structured result from the workflow
-                "messages": output.get("messages", []),
+            if self.structured_result is not None:
-                "final_message": self.final_message,
+                print("Using pre-computed structured result from workflow")
-                "final_model_used": self.final_model_used or output.get("current_model", "unknown"),
+                structured_message = self.structured_result
-                "user_question": user_query
+            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]}...")
+            # final_response = {
-            return 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": [],
+            "message": "Sorry i encountered an error while processing your request. Please try again.",
-                "final_message": "Sorry, I encountered an error while processing your visualization request. Please try again.",
+            "options": [],
-                "final_model_used": "error",
+            "requires_selection": False,
-                "user_question": user_query,
+            "end": "error",
-                "error": str(e)
+            "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"])
     

test.py

@@ -1,7 +1,6 @@
-# terminal_chat.py
 from src.llm.orchestrator import DroneBot, Message  # Adjust import path as needed
 import json
-def terminal_chat():
+async def terminal_chat():
     print("🚁 DroneBot Terminal Chat")
     print("Type 'exit' to quit.\n")
 
@@ -31,15 +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
-        response_json = json.loads(response["final_message"])
+        
-        print(f"🤖 DroneBot: {response_json}\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")