Update app.py

app.py

@@ -1,7 +1,6 @@
-import os
 import logging
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel
 from langchain_community.llms import Ollama
 from langchain.callbacks.manager import CallbackManager
@@ -11,6 +10,16 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 app = FastAPI()
+
+# Allow cross-origin requests (for browser-based JS)
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # Adjust for production if needed
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
 MODEL_NAME = 'tinyllama'
 
 def get_llm():
@@ -20,14 +29,36 @@ def get_llm():
 class Question(BaseModel):
     text: str
 
+class QueryWithContext(BaseModel):
+    question: str
+    context: str
+
 @app.get("/")
 def read_root():
     return {"Hello": f"Welcome to {MODEL_NAME} FastAPI"}
 
+@app.post("/ask")
+def ask(query: QueryWithContext):
+    prompt = f"""You are an expert in quantitative methods. Based on the following lecture notes, answer the user's question.
+
+Lecture notes:
+{query.context}
+
+User's question:
+{query.question}
+"""
+    llm = get_llm()
+    try:
+        response = llm.invoke(prompt)
+        return {"answer": response}
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
 @app.on_event("startup")
 async def startup_event():
     logger.info(f"Starting up with model: {MODEL_NAME}")
 
 @app.on_event("shutdown")
 async def shutdown_event():
-    logger.info("Shutting down")
+    logger.info("Shutting down")
+