Spaces:
Runtime error
Runtime error
File size: 2,170 Bytes
60956da da92430 2d02bb1 60956da 2d02bb1 60956da |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 |
import os
from uuid import uuid4
from fastapi import FastAPI
from fastapi.responses import FileResponse
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from pydantic import BaseModel, Field
from sentence_transformers import SentenceTransformer
import weaviate
from weaviate.classes.init import Auth
os.environ["HF_HOME"] = "/tmp/.cache/huggingface"
os.environ["TRANSFORMERS_CACHE"] = "/tmp/.cache/transformers"
os.environ["SENTENCE_TRANSFORMERS_CACHE"] = "/tmp/.cache/sentence_transformers"
WEAVIATE_URL = os.getenv("WEAVIATE_URL", "https://hrdhwtqlrkqmc8sfizwvpq.c0.asia-southeast1.gcp.weaviate.cloud")
WEAVIATE_API_KEY = os.getenv("WEAVIATE_API_KEY", "pMDX7ysJPkSTUMdV3gEwxhGmyB7wB301fLaJ")
CLASS_NAME = "PdfChunk"
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_methods=["*"],
allow_headers=["*"],
)
app.mount("/static", StaticFiles(directory="static"), name="static")
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
client = weaviate.connect_to_weaviate_cloud(
cluster_url=WEAVIATE_URL,
auth_credentials=Auth.api_key(WEAVIATE_API_KEY),
)
class QueryRequest(BaseModel):
text: str = Field(..., min_length=1)
top_k: int = Field(default=7, ge=1, le=100)
@app.get("/")
def root():
return {"message": "PDF RAG API is running. Use POST / to query."}
@app.get("/openapi.yaml")
def serve_openapi_yaml():
return FileResponse("static/openapi.yaml", media_type="text/yaml")
@app.post("/")
def query_weaviate(q: QueryRequest):
try:
query_vector = model.encode(q.text).tolist()
collection = client.collections.get(CLASS_NAME)
results = collection.query.near_vector(
near_vector=query_vector,
limit=q.top_k,
)
return {
"query": q.text,
"results": [
{
"text": obj.properties.get("text", ""),
"source": obj.properties.get("source", ""),
}
for obj in results.objects
]
}
except Exception as e:
return {"error": str(e)} |