ably.do/ollama_service.py

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)