Update app.py

app.py

@@ -127,10 +127,11 @@ def generate_policy_queries(claim_info_json: str) -> str:
     """
     global llm_model
     prompt = f"""
-    Analyze the following auto insurance claim to identify 3-5 key policy sections to consult:
-    - Focus on collision coverage, liability, deductibles, and relevant exclusions or endorsements.
+    Analyze the following auto insurance claim and generate exactly 2 short, keyword-based search queries to find the right policy sections.
+    - Example good queries: "collision coverage", "deductible limits", "exclusions".
+    - DO NOT write full sentences.
     - Claim Data: {claim_info_json}
-    - Return a JSON object with a 'queries' field containing a list of strings, e.g., {{"queries": ["query1", "query2"]}}. Do not include metadata fields.
+    - Return a JSON object strictly matching this schema: {{"queries": ["keyword 1", "keyword 2"]}}
     """
     try:
         messages = [{"role": "user", "content": prompt}]
@@ -154,7 +155,7 @@ def retrieve_policy_text(queries_json: str) -> str:
         policy_texts = []
         for q in query_strings:
             query_embedding = embedder.encode([q])[0].tolist()
-            results = collection.query(query_embeddings=[query_embedding], n_results=2)
+            results = collection.query(query_embeddings=[query_embedding], n_results=1)
             if results['documents']:
                 policy_texts.extend(results['documents'][0])
         return "\n\n".join(set(policy_texts)) # Set removes duplicate chunks
@@ -194,7 +195,7 @@ def generate_recommendation(claim_info_json: str, policy_text: str) -> str:
         return f"Error generating recommendation: {str(e)}"
 
 @tool
-def finalize_decision(claim_info_json: str, recommendation_json: str) -> str:
+def finalize_decision(claim_info_json: str | dict, recommendation_json: str | dict) -> str:
     """Finalize the claim decision based on the recommendation and format output.
     
     Args:
@@ -202,8 +203,16 @@ def finalize_decision(claim_info_json: str, recommendation_json: str) -> str:
         recommendation_json: The AI-generated recommendation output.
     """
     try:
-        claim_info = ClaimInfo.model_validate_json(claim_info_json)
-        rec_data = json.loads(recommendation_json)
+        # Gracefully handle if the LLM passes a dict OR a string
+        if isinstance(claim_info_json, dict):
+            claim_info = ClaimInfo.model_validate(claim_info_json)
+        else:
+            claim_info = ClaimInfo.model_validate_json(claim_info_json)
+            
+        if isinstance(recommendation_json, dict):
+            rec_data = recommendation_json
+        else:
+            rec_data = json.loads(recommendation_json)
         
         # Safe defaults if the model missed a key
         covered = "covered" in rec_data.get("recommendation_summary", "").lower() or (rec_data.get("settlement_amount", 0) or 0) > 0