Update app.py

app.py

@@ -6,6 +6,10 @@ from fastapi import FastAPI, Request, HTTPException
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 
+# --- New imports for LangChain + langchain_ollama ---
+from pydantic import BaseModel
+from langchain_ollama import Ollama
+
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
@@ -27,15 +31,25 @@ app.add_middleware(
 
 # Load your API key from the environment (defaults to "change_me")
 API_KEY = os.environ.get("API_KEY", "change_me")
-
 logger.info(f"API key loaded: {API_KEY}")
 
 # URL of the running Ollama server (adjust as needed)
 OLLAMA_SERVER_URL = "http://localhost:11434/api/generate"
 
+# --- Initialize a single Ollama instance via langchain_ollama ---
+#    This is the "LangChain" style interface to the Ollama server.
+ollama_llm = Ollama(
+    model="hf.co/abanm/Dubs-Q8_0-GGUF:latest",  # same model as before
+    base_url="http://localhost:11434",          # base URL for Ollama server
+    request_timeout=120                         # Increase if model loading is slow
+)
+
+# --------------------------------
+# Original endpoint: /api/generate
+# --------------------------------
 @app.post("/api/generate")
 async def generate(request: Request):
-    """Endpoint that generates text based on the prompt."""
+    """Endpoint that generates text based on the prompt (direct HTTP call to Ollama)."""
     try:
         # 1. Parse the incoming request
         body = await request.json()
@@ -74,7 +88,7 @@ async def generate(request: Request):
                         response.raise_for_status()
                         async for chunk in response.aiter_text():
                             yield chunk
-            except httpx.RequestError as exc:
+            except httpx.RequestError:
                 logger.exception("Request error while communicating with Ollama")
                 yield json.dumps({"error": "Unable to communicate with Ollama"})
             except httpx.HTTPStatusError as exc:
@@ -87,6 +101,38 @@ async def generate(request: Request):
         logger.exception("Unhandled exception in /api/generate")
         raise HTTPException(status_code=500, detail="Internal server error")
 
+# ------------------------------
+# New endpoint: /api/langchain-generate
+# ------------------------------
+class LangChainRequest(BaseModel):
+    prompt: str
+
+@app.post("/api/langchain-generate")
+async def langchain_generate(request: LangChainRequest):
+    """
+    Endpoint that uses langchain_ollama to generate text.
+    This is an alternative approach that uses the Ollama() class from langchain_ollama.
+    """
+    # Check for API key (similar logic as above, or unify the code)
+    # ... or just omit it if your environment is already secure
+    # If re-using the same approach:
+    #  (In real code, you'd unify these checks in a shared function)
+    #  
+    # auth_header = ...
+    # if token != API_KEY:
+    #    raise HTTPException(...)
+    
+    prompt = request.prompt
+    logger.info(f"LangChain request: {prompt}")
+
+    # Directly call the Ollama LLM via langchain_ollama
+    try:
+        response_text = ollama_llm(prompt)
+        return {"response": response_text}
+    except Exception as e:
+        logger.exception("Unhandled exception in /api/langchain-generate")
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.get("/health")
 async def health():
     """Health check endpoint."""
@@ -97,4 +143,3 @@ if __name__ == "__main__":
     import uvicorn
     logger.info("Starting server on http://0.0.0.0:7860")
     uvicorn.run(app, host="0.0.0.0", port=7860)
-