Files
ai.bewerbung/rag-api/main.py
T

152 lines
4.2 KiB
Python

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import StreamingResponse
from pydantic import BaseModel
from sentence_transformers import SentenceTransformer
import psycopg2
import httpx
import json
import os
app = FastAPI()
app.add_middleware(
CORSMiddleware,
allow_origins=["https://apply4jobs.de", "https://www.apply4jobs.de", "http://localhost:3000"],
allow_methods=["POST", "GET", "OPTIONS"],
allow_headers=["*"],
)
model = SentenceTransformer("intfloat/multilingual-e5-small")
DB_CONFIG = {
"host": "postgres",
"port": 5432,
"dbname": "anythingllm",
"user": "anythingllm",
"password": "any0203thing78llm"
}
OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", "")
OPENAI_MODEL = "gpt-4o-mini"
SYSTEM_PROMPT = """Du bist ein professioneller Assistent für Sebastian Fröhlich.
Beantworte ausschließlich Fragen zu seiner Person, seinen Fähigkeiten,
Projekten und Berufserfahrung. Nutze nur die bereitgestellten Dokumente.
Antworte auf Deutsch oder Englisch je nach Sprache des Recruiters.
Wenn du eine Frage nicht aus den Dokumenten beantworten kannst, sage das ehrlich.
Antworte immer in vollständigen, professionellen Sätzen."""
def get_context(query: str, top_k: int = 5) -> str:
query_text = f"query: {query}"
embedding = model.encode(query_text).tolist()
conn = psycopg2.connect(**DB_CONFIG)
cur = conn.cursor()
cur.execute("""
SELECT metadata->>'text'
FROM anythingllm_vectors
WHERE namespace = 'mein-workspace'
ORDER BY embedding <=> %s::vector
LIMIT %s
""", (embedding, top_k))
rows = cur.fetchall()
cur.close()
conn.close()
return "\n\n---\n\n".join([row[0] for row in rows])
class Message(BaseModel):
role: str
content: str
class ChatRequest(BaseModel):
model: str = OPENAI_MODEL
messages: list[Message]
stream: bool = False
class QueryRequest(BaseModel):
query: str
top_k: int = 5
@app.get("/health")
def health():
return {"status": "ok"}
@app.get("/v1/models")
def list_models():
return {
"object": "list",
"data": [{
"id": "sebastian-rag",
"object": "model",
"created": 1700000000,
"owned_by": "apply4jobs"
}]
}
@app.post("/retrieve")
def retrieve(req: QueryRequest):
context = get_context(req.query, req.top_k)
return {"context": context, "system_prompt": SYSTEM_PROMPT}
@app.post("/v1/chat/completions")
async def chat(req: ChatRequest):
user_message = next(
(m.content for m in reversed(req.messages) if m.role == "user"), ""
)
context = get_context(user_message)
enriched_messages = [
{
"role": "system",
"content": f"{SYSTEM_PROMPT}\n\n## Relevante Dokumente:\n{context}"
},
*[{"role": m.role, "content": m.content} for m in req.messages]
]
async def stream_response():
async with httpx.AsyncClient(timeout=60) as client:
async with client.stream(
"POST",
"https://api.openai.com/v1/chat/completions",
headers={
"Authorization": f"Bearer {OPENAI_API_KEY}",
"Content-Type": "application/json"
},
json={
"model": OPENAI_MODEL,
"messages": enriched_messages,
"stream": True
}
) as response:
async for chunk in response.aiter_bytes():
yield chunk
if req.stream:
return StreamingResponse(
stream_response(),
media_type="text/event-stream"
)
else:
async with httpx.AsyncClient(timeout=60) as client:
response = await client.post(
"https://api.openai.com/v1/chat/completions",
headers={
"Authorization": f"Bearer {OPENAI_API_KEY}",
"Content-Type": "application/json"
},
json={
"model": OPENAI_MODEL,
"messages": enriched_messages,
"stream": False
}
)
return response.json()