Update main.py

main.py

@@ -1,37 +1,55 @@
-from fastapi import FastAPI
-from duckduckgo_search import DDGS
-from fastapi.responses import JSONResponse
+from fastapi import FastAPI, HTTPException
 from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from duckai import DuckAI
 
 app = FastAPI()
-origins = ["*"] 
 
+# Configure CORS
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=origins,
+    allow_origins=["*"],
     allow_credentials=True,
-    allow_methods=["*"], 
-    allow_headers=["*"], 
+    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
+    allow_headers=["*"],
 )
 
-def chat_with_model(query: str, model: str) -> JSONResponse:
-    results = None
-    try:
-        results = DDGS().chat(query, model=model)
-    except Exception as e:
-        return JSONResponse(content={"error": str(e)}, status_code=500)
-    return JSONResponse(content={"results": results})
+class ChatQuery(BaseModel):
+    query: str
 
 @app.get("/chat/")
-async def chat(query: str) -> JSONResponse:
-    results = None
+async def chat(query: str):
+    if not query:
+        raise HTTPException(status_code=400, detail="Query parameter is required")
+
+    duck = DuckAI()
     try:
-        return  DDGS().chat(query, model='gpt-4o-mini')
-    except Exception as e:
+        results = duck.chat(query, model='gpt-4o-mini')
+        return {"results": results}
+    except Exception as e1:
+        print(f"Primary model (gpt-4o-mini) failed: {e1}")
         try:
-            return  DDGS().chat(query, model='claude-3-haiku')
-        except Exception as e:
-            return JSONResponse(content={"error": str(e)})
+            results = duck.chat(query, model='claude-3-haiku')
+            return {"results": results}
+        except Exception as e2:
+            print(f"Fallback model (claude-3-haiku) also failed: {e2}")
+            raise HTTPException(
+                status_code=500,
+                detail={
+                    "error": "Both models failed",
+                    "primary_error": str(e1),
+                    "fallback_error": str(e2)
+                }
+            )
+
+async def chat_with_model(query: str, model: str):
+    try:
+        duck = DuckAI()
+        results = duck.chat(query, model=model)
+        return {"results": results}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 
 if __name__ == "__main__":
     import uvicorn