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 import os app = FastAPI() OLLAMA_BASE_URL = "http://ollama:11434" WEAVIATE_URL = "http://weaviate:8080" PROMPT_DIR_PATCH = "./prompts" # 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") files_content = None 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 def read_text_files_from_directory(directory_path): files_dict = {} # Iterowanie przez wszystkie pliki w katalogu for filename in os.listdir(directory_path): # Sprawdzanie, czy plik ma rozszerzenie .txt if filename.endswith('.txt'): file_path = os.path.join(directory_path, filename) try: # Odczytywanie zawartości pliku with open(file_path, 'r', encoding='utf-8') as file: content = file.read() # Dodawanie do słownika z nazwą pliku bez rozszerzenia jako kluczem file_name_without_extension = os.path.splitext(filename)[0] files_dict[file_name_without_extension] = content except Exception as e: print(f"Błąd przy odczycie pliku {filename}: {e}") return files_dict def analyze_query(content): analysis = ollama_client.chat( model="gemma2:2b", messages=[{"role": "user", "content": content}] ) 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=100, 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: files_content = read_text_files_from_directory(PROMPT_DIR_PATCH) if files_content is None: raise KeyError(f"Nie wczytano promptów!!!") prompt_seach = files_content.get("prompt_seach") if prompt_seach is None: raise KeyError(f"Nie znaleziono promptu o nazwie '{prompt_seach}'.") prompt_system = files_content.get("prompt_system") if prompt_system is None: raise KeyError(f"Nie znaleziono promptu o nazwie '{prompt_system}'.") prompt_answer = files_content.get("prompt_answer") if prompt_answer is None: raise KeyError(f"Nie znaleziono promptu o nazwie '{prompt_answer}'.") prompt_data = files_content.get("prompt_data") if prompt_data is None: raise KeyError(f"Nie znaleziono promptu o nazwie '{prompt_data}'.") query = request.messages[-1].content if request.messages else "" keywords = analyze_query(prompt_seach.format(query=query)) weaviate_results = hybrid_search(keywords) prompt_data = "\n".join([f"Źródło: {doc['file_name']}\n{doc['relevant_fragment']}\n\n" for doc in weaviate_results]) print(prompt_data) messages_with_context =[ {"role": "system", "content": prompt_system}, {"role": "system", "content": prompt_data}, {"role": "user", "content": prompt_answer.format(query=query)} ] 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)