Update main.py

main.py

@@ -1,37 +1,51 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-import requests, uvicorn
+import requests
+import os
 
 app = FastAPI()
 
-# --- Gemini API Setup ---
-GEMINI_API_KEY = "AIzaSyApC_u1SScesLRbTfdz5lEcQExDaX-5IwQ"
-GEMINI_URL = f"https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent?key={GEMINI_API_KEY}"
+# --- Gemini API key ---
+GEMINI_API_KEY = "AIzaSyApC_u1SScesLRbTfdz5lEcQExDaX-5IwQ"  # Replace this with your real key
+GEMINI_URL = (
+    "https://generativelanguage.googleapis.com/v1beta/models/gemini-2.0-flash:generateContent"
+)
 
-# --- Input Model ---
+# --- Static prompt ---
+STATIC_INSTRUCTION = (
+    "Reply with a short, well-written summary in 2 to 3 sentences. "
+    "Be clear and concise. Avoid unnecessary explanations."
+)
+
+# --- Input schema ---
 class ChatRequest(BaseModel):
-    query: str   # Only query is sent by user
+    query: str
+
+# --- Output schema (optional) ---
+class ChatResponse(BaseModel):
+    reply: str
 
-# --- POST /chat ---
-@app.post("/chat")
+# --- Endpoint ---
+@app.post("/chat", response_model=ChatResponse)
 def chat(req: ChatRequest):
     try:
-        # Static instruction (used every time)
-        static_prompt = (
-            "You are an assistant. "
-            "Reply with a short, well-written summary in 2 to 3 sentences. "
-            "Do not explain unnecessarily. Be clear and concise."
-        )
+        # Combine prompt and query as a single user message
+        full_prompt = f"{STATIC_INSTRUCTION}\n\nUser question: {req.query}"
 
         contents = [
-            {"role": "system", "parts": [{"text": static_prompt}]},
-            {"role": "user",   "parts": [{"text": req.query}]}
+            {"role": "user", "parts": [{"text": full_prompt}]}
         ]
 
-        response = requests.post(GEMINI_URL,
-                                 headers={"Content-Type": "application/json"},
-                                 json={"contents": contents})
-        response.raise_for_status()
+        response = requests.post(
+            GEMINI_URL,
+            headers={
+                "Content-Type": "application/json",
+                "x-goog-api-key": GEMINI_API_KEY
+            },
+            json={"contents": contents}
+        )
+
+        response.raise_for_status()  # raise exception if status != 200
 
         data = response.json()
         reply = (
@@ -40,11 +54,7 @@ def chat(req: ChatRequest):
                 .get("text", "No reply")
         )
 
-        return {"reply": reply}
+        return {"reply": reply.strip()}
 
     except Exception as e:
         return {"reply": f"Error: {e}"}
-
-# --- Run Server ---
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)