updated bot chat api

api/models/responses.py

@@ -3,9 +3,7 @@ from pydantic import BaseModel
 from typing import Dict
 class ChatResponse(BaseModel):
     status: str
-    message: str
+    message: Dict
-   
-
    
 class SurveyAgentResponse(BaseModel):
     status: str

api/routes/chat.py

@@ -25,13 +25,12 @@ async def chat_ai(
         
         # Initialize DroneBot with history
         bot = DroneBot(history=history, use_openai_as_fallback=True)
-        
-        # Get response
         result = bot.chat(request.query)
-
+        message = json.loads(result["final_message"])
+        print(result)
         return ChatResponse(
             status="success",
-            message=result["final_message"]
+            message=message
         )
     except Exception as e:
         raise HTTPException(

src/prompts/templates/chat_templates.py

@@ -436,4 +436,5 @@ NOTE: THIS IS SOLELY YOUR TASK PLEASE, NO OTHER THING
     : DO NOT FABRICATE AVAILABLE SLOTS, ALWAYS CHECK FIRST
     : ALWAYS RESPOND IN VALID JSON FORMAT WITH THE REQUIRED FIELDS
     : RETURN ONLY JSON FORMAT NO EXPLANATION OR EXTRA INFORMATION BEFORE OR AFTER, DO NOT ADD ```json before or after
+    : 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
 """

test.py

@@ -0,0 +1,90 @@
+import os
+import json
+import asyncio
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import HumanMessage, SystemMessage
+from src.prompts.templates.flght_prompt import flight_prompt
+from src.config.llm_config import LlmConfig
+from config import Config
+from src.components.data_extraction.weather_data import DroneWeatherDataExtractor
+from logger import logger 
+
+
+class DroneAssessmentAgent:
+    def __init__(self):
+        self.llm = ChatOpenAI(
+            api_key=Config.OPENAI_API_KEY,
+            model=LlmConfig.openai.models.gpt_4o,
+            temperature=0.3
+        )
+        self.weather_extractor = DroneWeatherDataExtractor()
+
+    async def run(self, booking_form: str) -> dict:
+        """
+        Run the drone environmental & safety assessment agent.
+
+        Args:
+            booking_form (str): Structured booking form input
+
+        Returns:
+            dict: AI-generated structured output
+        """
+        logger.info("Starting DroneAssessmentAgent run...")
+
+        try:
+            # Step 1: Fetch weather data
+            weather_data = await self.weather_extractor.extract_booking_weather_data(booking_form)
+            booking_form_text = json.dumps(booking_form)
+
+            # Step 2: Generate prompt
+            prompt = flight_prompt(booking_form_text, weather_data)
+            messages = [
+                SystemMessage(content=prompt),
+                HumanMessage(content=booking_form_text)
+            ]
+
+            # Step 3: Invoke LLM
+            logger.debug("Sending prompt to LLM...")
+            response = self.llm.invoke(messages)
+
+            if hasattr(response, "content"):
+                response_text = response.content.strip()
+                logger.info("Received LLM response.")
+                logger.debug(f"LLM Raw Output (first 300 chars):\n{response_text[:300]}")
+
+                # Extract JSON block
+                start_idx = response_text.find('{')
+                end_idx = response_text.rfind('}') + 1
+
+                if start_idx == -1 or end_idx == -1:
+                    raise ValueError("No JSON object found in output")
+
+                json_str = response_text[start_idx:end_idx]
+                return json.loads(json_str)
+
+            else:
+                raise ValueError("LLM returned no usable content")
+
+        except Exception as e:
+            logger.exception("Error in DroneAssessmentAgent")
+            return {
+                "error": str(e),
+                "raw_response": response.content if 'response' in locals() and hasattr(response, "content") else None
+            }
+
+
+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())

test2.py

@@ -0,0 +1,32 @@
+booking_form_input = {
+        "job_id": "1043",
+        "job_overview": {
+            "job_number": "Job #1043",
+            "site_name": "Hightower Solar Farm",
+            "assigned_date": "January 23, 2025",
+            "status": "Scheduled"
+        },
+        "site_information": {
+            "site_name": "Hightower Solar Farm",
+            "region": "North England",
+            "full_address": "Grange Lane, Manchester M34 7TF",
+            "gps_coordinates": {
+                "latitude": "53.4408° N",
+                "longitude": "2.2426° W"
+            }
+        },
+        "timing": {
+            "start_time": "09:00 AM",
+            "end_time": "10:30 AM",
+            "survey_duration": "30–45 mins",
+            "buffer_time": "45 mins"
+        },
+        "form": {
+            "asset_type": "Solar Farm",
+            "system_size": "5.2 MW capacity, approximately 16,000 panels across 12 hectares",
+            "survey_purpose": "Insurance assessment",
+            "assigned_engineer": "David Wilson - 0161-555-0876",
+            "contact_person": "Sarah Thompson",
+            "contact_phone": "0161-555-0234"
+        }
+    }