Update app.py

app.py

@@ -18,7 +18,7 @@ import torch
 import pandas as pd
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file  # Import Safetensors loader
-from typing import Dict, Optional
+from typing import List, Dict, Optional
 
 # Initialize FastAPI app
 app = FastAPI()
@@ -40,6 +40,22 @@ class QueryRequest(BaseModel):
     query: str
     language_code: int = 0
 
+class MedicalProfile(BaseModel):
+    chronic_conditions: List[str]
+    symptoms: List[str]
+    food_restrictions: List[str]
+    mental_conditions: List[str]
+    daily_symptoms: List[str]
+
+class ChatQuery(BaseModel):
+    query: str
+    conversation_id: str
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+    timestamp: str
+
 def init_nltk():
     """Initialize NLTK resources"""
     try:
@@ -258,69 +274,89 @@ async def health_check():
     }
     return status
 
-@app.post("/api/query")
-async def process_query(request: QueryRequest):
-    """Main query processing endpoint"""
+@app.post("/api/chat")
+async def chat_endpoint(chat_query: ChatQuery):
     try:
-        query_text = request.query
-        language_code = request.language_code
-
-        if not models or not data.get('embeddings'):
-            raise HTTPException(
-                status_code=503,
-                detail="The system is currently initializing. Please try again in a few minutes."
-            )
-
-        try:
-            if language_code == 0:
-                query_text = translate_text(query_text, 'ar_to_en')
-
-            query_embedding = models['embedding'].encode([query_text])  # Ensure the model has been loaded and is available.
-            relevant_docs = query_embeddings(query_embedding)
-
-            if not relevant_docs:
-                return {
-                    'answer': 'No relevant information found. Please try a different query.',
-                    'success': True
-                }
-
-            doc_texts = [retrieve_document_text(doc_id) for doc_id, _ in relevant_docs]
-            doc_texts = [text for text in doc_texts if text.strip()]
-
-            if not doc_texts:
-                return {
-                    'answer': 'Unable to retrieve relevant documents. Please try again.',
-                    'success': True
-                }
-
-            rerank_scores = rerank_documents(query_text, doc_texts)
-            ranked_texts = [text for _, text in sorted(zip(rerank_scores, doc_texts), reverse=True)]
-
-            context = " ".join(ranked_texts[:3])
-            answer = generate_answer(query_text, context)
-
-            if language_code == 0:
-                answer = translate_text(answer, 'en_to_ar')
-
-            return {
-                'answer': answer,
-                'reranked_documents': ranked_texts,
-                'success': True
-            }
+        query_text = chat_query.query
+        query_embedding = models['embedding'].encode([query_text])
+        relevant_docs = query_embeddings(query_embedding)
+        
+        doc_texts = [retrieve_document_text(doc_id) for doc_id, _ in relevant_docs]
+        doc_texts = [text for text in doc_texts if text.strip()]
+        
+        rerank_scores = rerank_documents(query_text, doc_texts)
+        ranked_texts = [text for _, text in sorted(zip(rerank_scores, doc_texts), reverse=True)]
+        
+        context = " ".join(ranked_texts[:3])
+        answer = generate_answer(query_text, context)
+        
+        return {
+            "response": answer,
+            "conversation_id": chat_query.conversation_id,
+            "success": True
+        }
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
-        except Exception as e:
-            print(f"Error processing query: {e}")
-            raise HTTPException(
-                status_code=500,
-                detail="An error occurred while processing your query"
-            )
+@app.post("/api/resources")
+async def resources_endpoint(profile: MedicalProfile):
+    try:
+        context = f"""
+        Medical conditions: {', '.join(profile.chronic_conditions)}
+        Current symptoms: {', '.join(profile.daily_symptoms)}
+        Restrictions: {', '.join(profile.food_restrictions)}
+        Mental health: {', '.join(profile.mental_conditions)}
+        """
+        
+        query_embedding = models['embedding'].encode([context])
+        relevant_docs = query_embeddings(query_embedding)
+        doc_texts = [retrieve_document_text(doc_id) for doc_id, _ in relevant_docs]
+        doc_texts = [text for text in doc_texts if text.strip()]
+        
+        rerank_scores = rerank_documents(context, doc_texts)
+        ranked_docs = sorted(zip(relevant_docs, rerank_scores, doc_texts), key=lambda x: x[1], reverse=True)
+        
+        resources = []
+        for (doc_id, _), score, text in ranked_docs[:10]:
+            doc_info = data['df'][data['df']['id'] == doc_id].iloc[0]
+            resources.append({
+                "id": doc_id,
+                "title": doc_info['title'],
+                "content": text[:200],
+                "score": float(score)
+            })
+        
+        return {"resources": resources, "success": True}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
+@app.post("/api/recipes")
+async def recipes_endpoint(profile: MedicalProfile):
+    try:
+        recipe_query = f"Recipes and meals suitable for someone with: {', '.join(profile.chronic_conditions + profile.food_restrictions)}"
+        
+        query_embedding = models['embedding'].encode([recipe_query])
+        relevant_docs = query_embeddings(query_embedding)
+        doc_texts = [retrieve_document_text(doc_id) for doc_id, _ in relevant_docs]
+        doc_texts = [text for text in doc_texts if text.strip()]
+        
+        rerank_scores = rerank_documents(recipe_query, doc_texts)
+        ranked_docs = sorted(zip(relevant_docs, rerank_scores, doc_texts), key=lambda x: x[1], reverse=True)
+        
+        recipes = []
+        for (doc_id, _), score, text in ranked_docs[:10]:
+            doc_info = data['df'][data['df']['id'] == doc_id].iloc[0]
+            if 'recipe' in text.lower() or 'meal' in text.lower():
+                recipes.append({
+                    "id": doc_id,
+                    "title": doc_info['title'],
+                    "content": text[:200],
+                    "score": float(score)
+                })
+        
+        return {"recipes": recipes[:5], "success": True}
     except Exception as e:
-        print(f"Error in process_query: {e}")
-        raise HTTPException(
-            status_code=500,
-            detail=str(e)
-        )
+        raise HTTPException(status_code=500, detail=str(e))
 
 # Initialize application
 print("Initializing application...")