235 lines
8.3 KiB
Python
235 lines
8.3 KiB
Python
from fastapi import FastAPI, HTTPException, Request
|
|
from fastapi.responses import StreamingResponse
|
|
from pydantic import BaseModel
|
|
import ollama
|
|
import weaviate
|
|
from weaviate.connect import ConnectionParams
|
|
from weaviate.collections.classes.filters import Filter
|
|
import re
|
|
import json
|
|
import uvicorn
|
|
import httpx
|
|
from typing import List, Optional
|
|
import asyncio
|
|
|
|
app = FastAPI()
|
|
|
|
OLLAMA_BASE_URL = "http://ollama:11434"
|
|
WEAVIATE_URL = "http://weaviate:8080"
|
|
|
|
# Inicjalizacja klientów
|
|
ollama_client = ollama.Client(host=OLLAMA_BASE_URL)
|
|
weaviate_client = weaviate.WeaviateClient(
|
|
connection_params=ConnectionParams.from_params(
|
|
http_host="weaviate",
|
|
http_port=8080,
|
|
http_secure=False,
|
|
grpc_host="weaviate",
|
|
grpc_port=50051,
|
|
grpc_secure=False,
|
|
)
|
|
)
|
|
weaviate_client.connect()
|
|
collection = weaviate_client.collections.get("Document")
|
|
|
|
class Message(BaseModel):
|
|
role: str
|
|
content: str
|
|
|
|
class ChatRequest(BaseModel):
|
|
model: str
|
|
messages: List[Message]
|
|
stream: Optional[bool] = False
|
|
options: Optional[dict] = None
|
|
|
|
class ChatResponse(BaseModel):
|
|
model: str
|
|
created_at: str
|
|
message: Message
|
|
done: bool
|
|
total_duration: int
|
|
load_duration: int
|
|
prompt_eval_count: int
|
|
prompt_eval_duration: int
|
|
eval_count: int
|
|
eval_duration: int
|
|
|
|
prompt = """
|
|
Jesteś precyzyjnym narzędziem do generowania słów kluczowych z zakresu BHP i prawa pracy. Twoje zadanie to podanie WYŁĄCZNIE najistotniejszych słów do wyszukiwania w bazie dokumentów prawnych.
|
|
|
|
Ścisłe zasady:
|
|
1. Jeśli zapytanie dotyczy konkretnego artykułu:
|
|
- Podaj TYLKO numer artykułu i nazwę kodeksu (np. "Art. 154, Kodeks pracy").
|
|
- NIE dodawaj żadnych innych słów.
|
|
2. Jeśli zapytanie nie dotyczy konkretnego artykułu:
|
|
- Podaj maksymalnie 3 najbardziej specyficzne terminy związane z zapytaniem.
|
|
- Unikaj ogólnych słów jak "praca", "pracownik", "pracodawca", chyba że są częścią specjalistycznego terminu.
|
|
3. Używaj wyłącznie terminów, które z pewnością występują w dokumentach prawnych lub specjalistycznych opracowaniach.
|
|
4. NIE dodawaj własnych interpretacji ani rozszerzeń zapytania.
|
|
|
|
Odpowiedz TYLKO listą słów kluczowych oddzielonych przecinkami, bez żadnych dodatkowych wyjaśnień czy komentarzy.
|
|
|
|
Zapytanie: '{query}'
|
|
"""
|
|
|
|
def analyze_query(query):
|
|
analysis = ollama_client.chat(
|
|
model="gemma2:2b",
|
|
messages=[{"role": "user", "content": prompt.format(query=query)}]
|
|
)
|
|
keywords = [word.strip() for word in analysis['message']['content'].split(',') if word.strip()]
|
|
print("Słowa kluczowe:", keywords)
|
|
return keywords
|
|
|
|
def extract_full_article(content, article_number):
|
|
pattern = rf"Art\.\s*{article_number}\..*?(?=Art\.\s*\d+\.|\Z)"
|
|
match = re.search(pattern, content, re.DOTALL)
|
|
if match:
|
|
return match.group(0).strip()
|
|
return None
|
|
|
|
def extract_relevant_fragment(content, query, context_size=100):
|
|
article_match = re.match(r"Art\.\s*(\d+)", query)
|
|
if article_match:
|
|
article_number = article_match.group(1)
|
|
full_article = extract_full_article(content, article_number)
|
|
if full_article:
|
|
return full_article
|
|
|
|
index = content.lower().find(query.lower())
|
|
if index != -1:
|
|
start = max(0, index - context_size)
|
|
end = min(len(content), index + len(query) + context_size)
|
|
return f"...{content[start:end]}..."
|
|
return content[:200] + "..."
|
|
|
|
def hybrid_search(keywords, limit=5, alpha=0.5):
|
|
if isinstance(keywords, str):
|
|
keywords = [keywords]
|
|
|
|
query = " ".join(keywords)
|
|
|
|
print(f"\nWyszukiwanie hybrydowe dla słowa kluczowego: '{query}'")
|
|
response = collection.query.hybrid(
|
|
query=query,
|
|
alpha=alpha,
|
|
limit=limit * 2
|
|
)
|
|
|
|
results = []
|
|
|
|
for obj in response.objects:
|
|
#print(f"UUID: {obj.uuid}")
|
|
relevant_fragment = extract_relevant_fragment(obj.properties['content'], query)
|
|
#print(f"Relewantny fragment:\n{relevant_fragment}")
|
|
#print(f"Nazwa pliku: {obj.properties['fileName']}")
|
|
#print("---")
|
|
# Zmieniamy warunek na 'any' zamiast 'all'
|
|
#if any(term.lower() in relevant_fragment.lower() for term in keywords):
|
|
results.append({
|
|
"uuid": obj.uuid,
|
|
"relevant_fragment": relevant_fragment,
|
|
"file_name": obj.properties['fileName'],
|
|
"keyword": query
|
|
})
|
|
print(f"Dodano do wyników: {obj.uuid}")
|
|
|
|
if len(results) >= limit:
|
|
break
|
|
return results[:limit]
|
|
|
|
@app.get("/api/tags")
|
|
async def tags_proxy():
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.get(f"{OLLAMA_BASE_URL}/api/tags")
|
|
return response.json()
|
|
|
|
@app.get("/api/version")
|
|
async def tags_proxy():
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.get(f"{OLLAMA_BASE_URL}/api/version")
|
|
return response.json()
|
|
|
|
@app.post("/api/generate")
|
|
async def generate_proxy(request: Request):
|
|
data = await request.json()
|
|
async with httpx.AsyncClient() as client:
|
|
response = await client.post(f"{OLLAMA_BASE_URL}/api/generate", json=data)
|
|
return response.json()
|
|
|
|
@app.get("/api/models")
|
|
async def list_models():
|
|
try:
|
|
models = ollama_client.list()
|
|
return {"models": [model['name'] for model in models['models']]}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
async def stream_chat(model, messages, options):
|
|
try:
|
|
# Użycie httpx do asynchronicznego pobrania danych od Ollamy
|
|
async with httpx.AsyncClient() as client:
|
|
async with client.stream(
|
|
"POST",
|
|
f"{OLLAMA_BASE_URL}/api/chat",
|
|
json={"model": model, "messages": messages, "stream": True, "options": options},
|
|
) as response:
|
|
async for line in response.aiter_lines():
|
|
yield line + "\n"
|
|
except Exception as e:
|
|
yield json.dumps({"error": str(e)}) + "\n"
|
|
|
|
@app.post("/api/chat")
|
|
async def chat_endpoint(request: ChatRequest):
|
|
try:
|
|
query = request.messages[-1].content if request.messages else ""
|
|
keywords = analyze_query(query)
|
|
weaviate_results = hybrid_search(keywords)
|
|
|
|
if not weaviate_results:
|
|
context = f"""
|
|
Nie znalazłem informacji na temat: {query}.
|
|
Proszę poinformuj użytkownika, że nie masz wystarczającej wiedzy, aby udzielić jednoznacznej odpowiedzi.
|
|
"""
|
|
else:
|
|
context = "Znalezione informacje:\n"
|
|
for item in weaviate_results:
|
|
context += f"Źródło: {item['file_name']}\nFragment: {item['relevant_fragment']}\n\n"
|
|
|
|
messages_with_context =[
|
|
{"role": "system", "content": context},
|
|
{"role": "user", "content": f"""
|
|
Na podstawie powyższych informacji, odpowiedz na pytanie: {query}.
|
|
Odwołaj się do konkretnych artykułów lub zacytuj fragmenty źródeł.
|
|
"""}
|
|
]
|
|
|
|
if request.stream:
|
|
return StreamingResponse(stream_chat(request.model, messages_with_context, request.options), media_type="application/json")
|
|
|
|
ollama_response = ollama_client.chat(
|
|
model=request.model,
|
|
messages=messages_with_context,
|
|
stream=False,
|
|
options=request.options
|
|
)
|
|
return ChatResponse(
|
|
model=request.model,
|
|
created_at=ollama_response.get('created_at', ''),
|
|
message=Message(
|
|
role=ollama_response['message']['role'],
|
|
content=ollama_response['message']['content']
|
|
),
|
|
done=ollama_response.get('done', True),
|
|
total_duration=ollama_response.get('total_duration', 0),
|
|
load_duration=ollama_response.get('load_duration', 0),
|
|
prompt_eval_count=ollama_response.get('prompt_eval_count', 0),
|
|
prompt_eval_duration=ollama_response.get('prompt_eval_duration', 0),
|
|
eval_count=ollama_response.get('eval_count', 0),
|
|
eval_duration=ollama_response.get('eval_duration', 0)
|
|
)
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
if __name__ == "__main__":
|
|
uvicorn.run(app, host="0.0.0.0", port=8000) |