Update app.py

app.py

@@ -505,35 +505,50 @@ def translate_en_to_ar(text):
 
 
 
-# Medical context prompt
-MEDICAL_PROMPT = """You are a medical doctor who provides accurate and reliable health information based on current medical knowledge. 
-Only answer medical questions and provide information from reliable healthcare sources. 
-If a question is not medical in nature, politely explain that you can only address health-related queries.
-Question: {question}
-Answer: """
-
-def generate_response(question, max_length=350):
-    tok = models['gen_tokenizer']
-    mod = models['gen_model']    
-    # Prepare prompt
-    full_prompt = MEDICAL_PROMPT.format(question=question)
-    
-    # Generate response
-    inputs = tok(full_prompt, return_tensors="pt", truncation=True)
-    outputs = mod.generate(
-        inputs.input_ids,
-        max_length=max_length,
-        num_beams=4,
-        temperature=0.7,
-        early_stopping=True
-    )    
-    response = tok.decode(outputs[0], skip_special_tokens=True)
-    # Extract only the answer part
-    answer = response.split("Answer: ")[-1].strip()
-    return answer
 
 
 
+def get_completion(prompt: str, model: str = "gryphe/mythomax-l2-13b:free") -> str:
+    api_key = os.environ.get('OPENROUTER_API_KEY')
+    if not api_key:
+        raise HTTPException(status_code=500, detail="OPENROUTER_API_KEY not found in environment variables")
+    
+    client = OpenAI(
+        base_url="https://openrouter.ai/api/v1",
+        api_key=api_key
+    )
+    
+    if not prompt.strip():
+        raise HTTPException(status_code=400, detail="Please enter a question")
+    
+    try:
+        completion = client.chat.completions.create(
+            extra_headers={
+                "HTTP-Referer": "https://huggingface.co/spaces/thechaiexperiment/phitrial",
+                "X-Title": "My Hugging Face Space"
+            },
+            model=model,
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ]
+        )
+        
+        if (completion and 
+            hasattr(completion, 'choices') and 
+            completion.choices and 
+            hasattr(completion.choices[0], 'message') and 
+            hasattr(completion.choices[0].message, 'content')):
+            return completion.choices[0].message.content
+        else:
+            raise HTTPException(status_code=500, detail="Received invalid response from API")
+            
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
 
 
 
@@ -565,21 +580,41 @@ async def health_check():
 
 
 
-
-@app.post("/api/ask")
-async def chat_endpoint(chat_query: ChatQuery):
+@app.post("/ask")
+async def chat(query: ChatQuery):
     try:
-        query_text = chat_query.query
-        language_code = chat_query.language_code        
-        if language_code == 0:
-            query_text = translate_ar_to_en(query_text)
-        # Generate response
-        answer = generate_response(query_text)        
-        # Translate back to Arabic if needed
-        if language_code == 0:
-            answer = translate_en_to_ar(answer)
-        
-        return answer
+        # Handle Arabic input
+        if query.language_code == 0:
+            # Translate question from Arabic to English
+            english_query = translate_ar_to_en(query.query)
+            if not english_query:
+                raise HTTPException(status_code=500, detail="Failed to translate question from Arabic to English")
+            
+            # Get completion in English
+            english_response = get_completion(english_query)
+            
+            # Translate response back to Arabic
+            arabic_response = translate_en_to_ar(english_response)
+            if not arabic_response:
+                raise HTTPException(status_code=500, detail="Failed to translate response to Arabic")
+                
+            return {
+                "original_query": query.query,
+                "translated_query": english_query,
+                "response": arabic_response,
+                "response_in_english": english_response
+            }
+            
+        # Handle English input
+        else:
+            response = get_completion(query.query)
+            return {
+                "query": query.query,
+                "response": response
+            }
+            
+    except HTTPException as e:
+        raise e
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
@@ -589,6 +624,8 @@ async def chat_endpoint(chat_query: ChatQuery):
 
 
 
+
+
     
 
 @app.post("/api/chat")