From 02852b2992de118403d65343166e3a36623dfa1c Mon Sep 17 00:00:00 2001 From: OwusuBlessing Date: Mon, 18 Aug 2025 21:21:08 +0100 Subject: [PATCH] fixed llm issues with response --- api/routes/chat.py | 15 +--- src/llm/agent/structured_output.py | 110 +++++++++++++++++++++++++++++ src/llm/orchestrator.py | 88 +++++++++++++++++------ test.py | 14 ++-- tests/integration/bot_chat.py | 2 +- 5 files changed, 188 insertions(+), 41 deletions(-) create mode 100644 src/llm/agent/structured_output.py diff --git a/api/routes/chat.py b/api/routes/chat.py index ba65e5b..1f5fd64 100644 --- a/api/routes/chat.py +++ b/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") diff --git a/src/llm/agent/structured_output.py b/src/llm/agent/structured_output.py new file mode 100644 index 0000000..813e79f --- /dev/null +++ b/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()) diff --git a/src/llm/orchestrator.py b/src/llm/orchestrator.py index 2b27b86..0e81649 100644 --- a/src/llm/orchestrator.py +++ b/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 = { - "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"]) \ No newline at end of file diff --git a/test.py b/test.py index 5fd9a49..6b55509 100644 --- a/test.py +++ b/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()) \ No newline at end of file diff --git a/tests/integration/bot_chat.py b/tests/integration/bot_chat.py index 49f5e73..2cc1b77 100644 --- a/tests/integration/bot_chat.py +++ b/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")