Create app.py

app.py

@@ -0,0 +1,66 @@
+from fastapi import FastAPI, UploadFile, File, Form
+from pydantic import BaseModel
+import openai
+import faiss
+import numpy as np
+import os
+from dotenv import load_dotenv
+from fastapi.middleware.cors import CORSMiddleware
+from PyPDF2 import PdfReader
+
+load_dotenv()
+openai.api_key = os.getenv("OPENAI_API_KEY")
+
+app = FastAPI()
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+notebooks = {}
+
+class Query(BaseModel):
+    question: str
+    notebook_id: str
+
+@app.post("/ask")
+def ask(query: Query):
+    nb = notebooks.get(query.notebook_id)
+    if not nb:
+        return {"answer": "Notebook not found."}
+    question_embedding = openai.Embedding.create(
+        input=[query.question],
+        model="text-embedding-ada-002"
+    )["data"][0]["embedding"]
+    if len(nb["texts"]) == 0:
+        return {"answer": "No documents indexed in this notebook."}
+    D, I = nb["index"].search(np.array([question_embedding]).astype("float32"), k=3)
+    context = "\n\n".join([f"[{i+1}] {nb['texts'][i]}" for i in I[0]])
+    citation_refs = [nb['citations'][i] for i in I[0]]
+    response = openai.ChatCompletion.create(
+        model="gpt-4",
+        messages=[
+            {"role": "system", "content": "You are an AI assistant that answers based on uploaded documents. Cite sources using [1], [2], etc."},
+            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {query.question}"}
+        ],
+        temperature=0.3
+    )
+    return {"answer": response.choices[0].message.content.strip(), "citations": citation_refs}
+
+@app.post("/upload-pdf")
+def upload_pdf(notebook_id: str = Form(...), file: UploadFile = File(...)):
+    if notebook_id not in notebooks:
+        notebooks[notebook_id] = {"index": faiss.IndexFlatL2(1536), "texts": [], "citations": []}
+    nb = notebooks[notebook_id]
+    reader = PdfReader(file.file)
+    for i, page in enumerate(reader.pages):
+        content = page.extract_text()
+        if content:
+            embedding = openai.Embedding.create(input=[content], model="text-embedding-ada-002")["data"][0]["embedding"]
+            nb["index"].add(np.array([embedding]).astype("float32"))
+            nb["texts"].append(content)
+            nb["citations"].append(f"{file.filename}, page {i+1}")
+    return {"status": f"{file.filename} uploaded and parsed"}