espiritismo-rag-system/backend/app/services/rag_service.py

from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct, Filter, FieldCondition, MatchValue
import numpy as np
from typing import List, Dict, Optional
import os
import json
import time

class QdrantService:
    def __init__(self, url: str, collection_name: str = "espiritismo_books"):
        self.client = QdrantClient(url=url, prefer_grpc=False)
        self.collection_name = collection_name
        self.embedding_dim = 384

        self._ensure_collection()

    def _ensure_collection(self):
        """Create collection if it doesn't exist."""
        last_error = None
        for attempt in range(15):
            try:
                collections = self.client.get_collections().collections
                break
            except Exception as exc:
                last_error = exc
                print(f"Qdrant indisponível ({attempt + 1}/15): {exc}")
                time.sleep(2)
        else:
            raise RuntimeError(f"Não foi possível conectar ao Qdrant: {last_error}")

        collection_exists = any(
            c.name == self.collection_name for c in collections
        )

        if not collection_exists:
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=self.embedding_dim,
                    distance=Distance.COSINE,
                ),
            )
            print(f"Created collection: {self.collection_name}")
        else:
            print(f"Collection {self.collection_name} already exists")

    def add_chunks(self, chunks: List[Dict]):
        """Add chunks to the collection."""
        points = []

        for i, chunk in enumerate(chunks):
            point = PointStruct(
                id=chunk.get("id") or i,
                vector=chunk["embedding"],
                payload={
                    "content": chunk["content"],
                    "metadata": chunk["metadata"],
                }
            )
            points.append(point)

        if points:
            self.client.upsert(
                collection_name=self.collection_name,
                points=points,
            )
            print(f"Added {len(points)} chunks to Qdrant")

    def search(self, query_embedding, top_k: int = 3) -> List[Dict]:
        """Search for similar chunks."""
        if hasattr(query_embedding, "tolist"):
            query_embedding = query_embedding.tolist()
        results = self.client.search(
            collection_name=self.collection_name,
            query_vector=query_embedding,
            limit=top_k,
        )

        return [
            {
                "content": hit.payload["content"],
                "metadata": hit.payload["metadata"],
                "score": hit.score,
            }
            for hit in results
        ]

    def get_books(self) -> List[str]:
        """Get list of books in the collection."""
        points = self.client.scroll(
            collection_name=self.collection_name,
            limit=10000,
        )[0]

        books = set()
        for point in points:
            metadata = point.payload["metadata"]
            books.add(metadata.get("title", "Unknown"))

        return sorted(list(books))

    def get_collection_info(self) -> Dict:
        """Get information about the collection."""
        info = self.client.get_collection(self.collection_name)
        return {
            "status": info.status,
            "vectors_count": info.vectors_count,
            "points_count": info.points_count,
        }


class RAGService:
    def __init__(self, qdrant_service: QdrantService, model_name: str = "llama-3.2-3b-instruct-q4_0.gguf"):
        self.qdrant = qdrant_service
        self.model_name = model_name
        self.llm = None

    def _get_llm(self):
        if self.llm is None:
            if not os.path.isfile(self.model_name):
                raise FileNotFoundError(
                    f"Modelo não encontrado em {self.model_name}. "
                    "Baixe llama-3.2-3b-instruct-q4_0.gguf para backend/models/."
                )
            from gpt4all import GPT4All

            model_dir = os.path.dirname(self.model_name) or "."
            model_file = os.path.basename(self.model_name)
            self.llm = GPT4All(
                model_name=model_file,
                model_path=model_dir,
                device="cpu",
                allow_download=False,
            )
        return self.llm

    def generate_response(self, query: str, context: List[str], top_k: int = 3) -> str:
        """Generate a response based on the query and context."""
        context_text = "\n\n".join(context)

        system_prompt = f"""Você é um assistente especializado em espiritismo. Use as informações abaixo para responder à pergunta do usuário.

Informações:
{context_text}

Instruções:
- Responda de forma clara e concisa
- Use o contexto fornecido
- Quando citar passagens, mencione o livro e capítulo (se disponível)
- Se não souber a resposta, diga que não encontrou na base de dados"""

        prompt = f"{system_prompt}\n\nPergunta: {query}\nResposta:"
        response = self._get_llm().generate(prompt, max_tokens=500)

        return response

    def query(self, query: str, top_k: int = 3) -> Dict:
        """Query the RAG system."""
        # Get embeddings
        from app.utils.embeddings import EmbeddingService
        embedding_service = EmbeddingService()

        query_embedding = embedding_service.embed_query(query)

        # Search
        chunks = self.qdrant.search(query_embedding, top_k)

        # Generate response
        context = [chunk["content"] for chunk in chunks]

        response = self.generate_response(query, context, top_k)

        return {
            "query": query,
            "response": response,
            "context": context,
            "sources": [
                {"metadata": chunk["metadata"], "score": chunk["score"]}
                for chunk in chunks
            ],
        }