Update main.py

main.py

@@ -1,4 +1,5 @@
 import os
+import json
 from datetime import datetime, timedelta
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
@@ -11,18 +12,6 @@ from langchain.prompts import SystemMessagePromptTemplate, HumanMessagePromptTem
 from langchain.prompts.chat import ChatPromptTemplate
 import pytz
 import uvicorn
-import json
-
-def safe_json_parse(content: str):
-    try:
-        while isinstance(content, str):
-            content = content.strip()
-            content = content.strip('"').strip("'")  # remove outer quotes
-            content = json.loads(content)  # keep parsing until it becomes dict
-        return content
-    except Exception as e:
-        raise ValueError(f"JSON parsing failed: {e}")
-
 
 # 🔐 Load environment
 load_dotenv()
@@ -51,7 +40,7 @@ def get_week_range():
     end = start + timedelta(days=6)
     return start.strftime("%Y-%m-%d"), end.strftime("%Y-%m-%d")
 
-# 🧾 SQL Builder with optional date range
+# 🧾 SQL Builder
 def build_sql(first_name: str, last_name: str = None, start_date: str = None, end_date: str = None) -> str:
     if not start_date or not end_date:
         start_date, end_date = get_week_range()
@@ -85,10 +74,6 @@ def build_sql(first_name: str, last_name: str = None, start_date: str = None, en
 # 🛠️ LangChain Tool
 @tool
 def query_weekly_progress(query: str) -> str:
-    """
-    Executes a SELECT query to get user progress for the week.
-    Accepts SQL generated from name and timeframe.
-    """
     try:
         print(f"📥 SQL: {query}")
         if not query.lower().strip().startswith("select"):
@@ -98,7 +83,7 @@ def query_weekly_progress(query: str) -> str:
     except Exception as e:
         return f"❌ SQL Error: {e}"
 
-# 🤖 Weekly report formatter (HTML)
+# 🤖 Weekly report formatter
 system_prompt = SystemMessagePromptTemplate(
     prompt=PromptTemplate(
         input_variables=[],
@@ -131,7 +116,7 @@ chat_prompt = ChatPromptTemplate.from_messages([
 agent = create_openai_tools_agent(llm=llm, tools=[query_weekly_progress], prompt=chat_prompt)
 agent_executor = AgentExecutor(agent=agent, tools=[query_weekly_progress])
 
-# ✨ Prompt: Extract name and optional date range
+# ✨ Name + Date extractor prompt
 name_extractor = ChatOpenAI(model="gpt-3.5-turbo", temperature=0, openai_api_key=OPENAI_API_KEY)
 name_prompt = PromptTemplate(
     input_variables=["text"],
@@ -152,20 +137,36 @@ Prompt: {text}
 """
 )
 
-# 📥 Request model
+# 📥 Input model
 class PromptRequest(BaseModel):
     prompt: str
 
-# 🔍 Endpoint
+# 🧠 Safe JSON parser for LLM responses
+def safe_extract_json(raw_content: str) -> dict:
+    try:
+        while isinstance(raw_content, str):
+            raw_content = raw_content.strip().strip('"').strip("'")
+            raw_content = json.loads(raw_content)
+        if isinstance(raw_content, dict):
+            return raw_content
+        raise ValueError("Parsed content is not a dict.")
+    except Exception as e:
+        raise ValueError(f"Could not extract JSON from model response: {e}")
+
+# 🔍 Main endpoint
 @app.post("/weekly-dpr")
 async def get_prompt_based_dpr(request: PromptRequest):
     try:
         prompt_text = request.prompt.strip()
-        name_response = name_extractor.invoke(name_prompt.format_prompt(text=prompt_text).to_string())
+        name_response = name_extractor.invoke(
+            name_prompt.format_prompt(text=prompt_text).to_string()
+        )
+
+        print("🧠 Raw model response:", repr(name_response.content))
 
-        print("🧠 Extracted JSON string:", name_response.content)
+        name_json = safe_extract_json(name_response.content)
 
-        name_json = safe_json_parse(name_response.content)
+        print("✅ Parsed JSON:", name_json)
 
         first = name_json.get("first_name")
         last = name_json.get("last_name")
@@ -173,7 +174,7 @@ async def get_prompt_based_dpr(request: PromptRequest):
         end_date = name_json.get("end_date")
 
         if not first:
-            raise ValueError("No valid name extracted.")
+            raise ValueError("No valid first name extracted.")
 
         sql = build_sql(first, last, start_date, end_date)
         context = f"Generate HTML weekly report for: {first} {last or ''} from {start_date or 'this week'} to {end_date or ''} using this query:\n{sql}"
@@ -184,7 +185,6 @@ async def get_prompt_based_dpr(request: PromptRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"DPR Error: {e}")
 
-
 # ▶️ Run the app
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)