152 lines
4.2 KiB
Python
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()
|