agent/tools/rag_tool.py

import json
import uuid
import chromadb
from crewai.tools import BaseTool
from pydantic import BaseModel, Field
import config

_chroma_client = chromadb.EphemeralClient()
_embedder = None


def _get_embedder():
    global _embedder
    if _embedder is None:
        from sentence_transformers import SentenceTransformer
        _embedder = SentenceTransformer("BAAI/bge-small-en-v1.5")
    return _embedder


def _safe_collection_name(session_id: str) -> str:
    clean = session_id.replace("-", "")[:40]
    return f"s{clean}"


def _chunk_text(text: str, size: int = config.CHUNK_SIZE, overlap: int = config.CHUNK_OVERLAP) -> list[str]:
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk = " ".join(words[i: i + size])
        chunks.append(chunk)
        i += size - overlap
        if i >= len(words):
            break
    return chunks


class EmbedInput(BaseModel):
    sources_json: str = Field(description="JSON list of accepted sources with content and metadata")
    session_id: str = Field(description="Unique session ID to scope the ChromaDB collection")


class RetrieveInput(BaseModel):
    query: str = Field(description="Query to retrieve relevant chunks for")
    session_id: str = Field(description="Session ID matching the embed step")
    n_results: int = Field(default=10, description="Number of chunks to retrieve")


class EmbedSourcesTool(BaseTool):
    name: str = "Embed Sources into Vector Store"
    description: str = (
        "Chunk and embed accepted sources into a session-scoped ChromaDB collection. "
        "Must be called before RetrieveChunksTool. Returns the session_id for retrieval."
    )
    args_schema: type[BaseModel] = EmbedInput

    def _run(self, sources_json: str, session_id: str) -> str:
        try:
            sources = json.loads(sources_json)
        except json.JSONDecodeError as e:
            return json.dumps({"error": f"Invalid JSON: {e}"})

        if not isinstance(sources, list):
            sources = sources.get("accepted", []) if isinstance(sources, dict) else []

        collection_name = _safe_collection_name(session_id)
        try:
            _chroma_client.delete_collection(collection_name)
        except Exception:
            pass
        collection = _chroma_client.create_collection(collection_name)

        total_chunks = 0
        for src in sources:
            content = src.get("content", "")
            if not content or len(content.strip()) < 50:
                continue
            chunks = _chunk_text(content)
            for i, chunk in enumerate(chunks):
                try:
                    embedding = _get_embedder().encode(chunk).tolist()
                    collection.add(
                        ids=[str(uuid.uuid4())],
                        embeddings=[embedding],
                        documents=[chunk],
                        metadatas=[{
                            "url": src.get("url", ""),
                            "title": src.get("title", ""),
                            "credibility_score": str(src.get("credibility_score", 0)),
                            "confidence": src.get("confidence", "low"),
                            "chunk_index": str(i),
                        }],
                    )
                    total_chunks += 1
                except Exception:
                    continue

        return json.dumps({"embedded_chunks": total_chunks, "session_id": session_id})


class RetrieveChunksTool(BaseTool):
    name: str = "Retrieve Relevant Chunks"
    description: str = (
        "Retrieve the most relevant document chunks from the vector store for a query. "
        "Call EmbedSourcesTool first. Returns chunks with source metadata for grounded synthesis."
    )
    args_schema: type[BaseModel] = RetrieveInput

    def _run(self, query: str, session_id: str, n_results: int = 10) -> str:
        collection_name = _safe_collection_name(session_id)
        try:
            collection = _chroma_client.get_collection(collection_name)
        except Exception as e:
            return json.dumps({"error": f"Collection not found: {e}", "chunks": []})

        count = collection.count()
        if count == 0:
            return json.dumps({"error": "Collection is empty — embed sources first", "chunks": []})

        query_embedding = _embedder.encode(query).tolist()
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=min(n_results, count),
        )

        chunks = []
        docs = results.get("documents", [[]])[0]
        metas = results.get("metadatas", [[]])[0]
        for doc, meta in zip(docs, metas):
            chunks.append({
                "text": doc,
                "url": meta.get("url", ""),
                "title": meta.get("title", ""),
                "credibility_score": float(meta.get("credibility_score", 0)),
                "confidence": meta.get("confidence", "low"),
            })

        return json.dumps({"chunks": chunks, "count": len(chunks)})