mirror of
https://github.com/PlatypusPus/MushroomEmpire.git
synced 2026-02-07 22:18:59 +00:00
74 lines
2.0 KiB
Python
74 lines
2.0 KiB
Python
import ollama
|
|
import chromadb
|
|
from pypdf import PdfReader
|
|
from fastapi import FastAPI, HTTPException
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel
|
|
import uvicorn
|
|
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
|
|
app = FastAPI()
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"], # or specify ["http://localhost:3000"]
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
|
|
client = chromadb.Client()
|
|
collection = client.create_collection(name="docs")
|
|
|
|
reader = PdfReader('../../GDPRArticles.pdf')
|
|
chunks = [page.extract_text().strip() for page in reader.pages if page.extract_text().strip()]
|
|
print("Done reading pages!")
|
|
for i, chunk in enumerate(chunks):
|
|
response = ollama.embed(model="nomic-embed-text", input=chunk)
|
|
embedding = response["embeddings"][0]
|
|
collection.add(
|
|
ids=[str(i)],
|
|
embeddings=[embedding],
|
|
documents=[chunk]
|
|
)
|
|
print("Embeddings done!")
|
|
|
|
|
|
# Allow browser calls from the frontend
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
class ChatRequest(BaseModel):
|
|
prompt: str
|
|
|
|
@app.post("/chat")
|
|
async def chat_bot(prompt: str | None = None, body: ChatRequest | None = None):
|
|
# Accept prompt from either query (?prompt=) or JSON body {"prompt": "..."}
|
|
query = prompt or (body.prompt if body else None)
|
|
if not query:
|
|
raise HTTPException(status_code=400, detail="Missing prompt")
|
|
response = ollama.embed(model="nomic-embed-text", input=query)
|
|
query_embedding = response["embeddings"][0]
|
|
|
|
results = collection.query(
|
|
query_embeddings=[query_embedding],
|
|
n_results=1
|
|
)
|
|
data = results['documents'][0][0]
|
|
|
|
output = ollama.generate(
|
|
model="llama3.2",
|
|
prompt=f"Context: {data}\n\nQuestion: {query}\n\nAnswer:"
|
|
)
|
|
return {"response": output["response"]}
|
|
|
|
|
|
if __name__ == '__main__':
|
|
uvicorn.run(app, port=8080, host='0.0.0.0') |