ff-aitoolkit 0.2.0__py3-none-any.whl

aitoolkit/rag/agent.py

@@ -0,0 +1,165 @@
+"""Centralized RAG agent coordinating embeddings, vector store, and retrieval.
+
+Provider-agnostic: no Google/Gemini defaults. ``answer_question`` is functional —
+it retrieves context and generates an answer with the toolkit's
+:class:`~aitoolkit.llm.LLMClient`.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from loguru import logger
+
+from aitoolkit.config import get_settings
+from aitoolkit.embeddings import EmbeddingsClient, get_embeddings_client
+from aitoolkit.llm import LLMClient, get_llm_client
+from aitoolkit.rag.retriever import FilterFn, RAGRetriever
+from aitoolkit.rag.vector_store import UnifiedVectorStore
+from aitoolkit.types import RetrievedChunk
+
+_DEFAULT_ANSWER_SYSTEM = (
+    "You are a helpful assistant. Answer the question using ONLY the provided "
+    "context. If the context is insufficient, say so clearly."
+)
+
+
+class RAGAgent:
+    """Unified interface for embedding, storing, retrieving and answering."""
+
+    def __init__(
+        self,
+        qdrant_url: Optional[str] = None,
+        collection_name: Optional[str] = None,
+        redis_url: Optional[str] = None,
+        embeddings: Optional[EmbeddingsClient] = None,
+        llm: Optional[LLMClient] = None,
+        enable_caching: bool = True,
+        filter_fn: Optional[FilterFn] = None,
+    ) -> None:
+        settings = get_settings()
+        self.collection_name = collection_name or settings.qdrant_collection
+        self.embeddings = embeddings or get_embeddings_client()
+        self.llm = llm or get_llm_client()
+
+        self.vector_store = UnifiedVectorStore(
+            qdrant_url=qdrant_url,
+            collection_name=self.collection_name,
+            embeddings=self.embeddings,
+        )
+        self.retriever = RAGRetriever(
+            vector_store=self.vector_store,
+            redis_url=(redis_url or settings.redis_url) if enable_caching else None,
+            cache_ttl=settings.cache_ttl,
+            filter_fn=filter_fn,
+        )
+        logger.success(f"RAGAgent ready (collection='{self.collection_name}')")
+
+    async def add_documents(
+        self,
+        texts: List[str],
+        *,
+        file_id: str,
+        metadatas: Optional[List[Dict[str, Any]]] = None,
+        source_type: str = "upload",
+        **extra: Any,
+    ) -> List[str]:
+        """Embed and store document chunks. Returns point IDs."""
+        return await self.vector_store.add_texts(
+            texts, metadatas, file_id=file_id, source_type=source_type, **extra
+        )
+
+    async def delete_file(self, file_id: str) -> None:
+        await self.vector_store.delete_by_file_id(file_id)
+
+    async def retrieve_context(
+        self,
+        query: str,
+        *,
+        file_ids: Optional[List[str]] = None,
+        limit: int = 10,
+        score_threshold: Optional[float] = None,
+        use_cache: bool = True,
+    ) -> List[RetrievedChunk]:
+        return await self.retriever.retrieve(
+            query,
+            file_ids=file_ids,
+            limit=limit,
+            score_threshold=score_threshold,
+            use_cache=use_cache,
+        )
+
+    async def get_formatted_context(
+        self,
+        query: str,
+        *,
+        file_ids: Optional[List[str]] = None,
+        limit: int = 10,
+        separator: str = "\n\n---\n\n",
+    ) -> str:
+        return await self.retriever.get_context_text(
+            query, file_ids=file_ids, limit=limit, separator=separator
+        )
+
+    async def answer_question(
+        self,
+        question: str,
+        *,
+        file_ids: Optional[List[str]] = None,
+        limit: int = 5,
+        system: str = _DEFAULT_ANSWER_SYSTEM,
+        temperature: Optional[float] = None,
+    ) -> Dict[str, Any]:
+        """Retrieve context and generate a grounded answer."""
+        chunks = await self.retrieve_context(
+            question, file_ids=file_ids, limit=limit
+        )
+        if not chunks:
+            return {
+                "answer": "I don't have enough context to answer this question.",
+                "sources": [],
+            }
+
+        context = "\n\n---\n\n".join(c.text for c in chunks)
+        prompt = (
+            f"Context:\n{context}\n\nQuestion: {question}\n\nAnswer:"
+        )
+        answer = await self.llm.chat(
+            prompt, system=system, temperature=temperature
+        )
+        return {
+            "answer": answer,
+            "sources": [c.model_dump() for c in chunks],
+        }
+
+    async def get_file_ids(self, exclude: Optional[List[str]] = None) -> List[str]:
+        return await self.vector_store.get_unique_file_ids(exclude=exclude)
+
+    async def get_file_count(self, exclude: Optional[List[str]] = None) -> int:
+        return len(await self.get_file_ids(exclude=exclude))
+
+    async def aclose(self) -> None:
+        await self.retriever.aclose()
+        await self.vector_store.aclose()
+
+
+_agents: Dict[str, RAGAgent] = {}
+
+
+def get_rag_agent(
+    collection_name: Optional[str] = None,
+    **kwargs: Any,
+) -> RAGAgent:
+    """Return a cached RAG agent, one per collection.
+
+    Caching is keyed on the collection name so an app that serves multiple
+    collections gets a distinct agent for each — a single global singleton would
+    silently hand back the first collection's agent for every later call.
+    ``kwargs`` only take effect the first time a given collection is requested.
+    """
+    key = collection_name or get_settings().qdrant_collection
+    agent = _agents.get(key)
+    if agent is None:
+        agent = RAGAgent(collection_name=collection_name, **kwargs)
+        _agents[key] = agent
+    return agent

aitoolkit/rag/query_expansion.py

@@ -0,0 +1,147 @@
+"""Query expansion for improved RAG recall.
+
+Domain-agnostic: no domain-specific keyword list is hardcoded. Pass ``domain``
+and an optional ``keywords`` list to tailor expansion to any project. With no
+keywords, expansion is purely structural.
+"""
+
+from __future__ import annotations
+
+from functools import lru_cache
+from typing import List, Optional, Tuple
+
+from loguru import logger
+
+# Generic English stop words for naive keyword extraction.
+_STOP_WORDS = {
+    "the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for",
+    "of", "with", "is", "are", "was", "were", "be", "this", "that", "it",
+}
+
+# Content-type templates that anticipate document content for better retrieval.
+_TEMPLATES = {
+    "quiz": (
+        "What are the key concepts, definitions, procedures, and important points "
+        "related to {topic}? What questions should learners understand about {topic}?"
+    ),
+    "lesson": (
+        "Explain {topic} in detail. What are the fundamental concepts of {topic}? "
+        "What should learners understand about {topic}?"
+    ),
+    "summary": (
+        "Summarize the main points about {topic}. What are the key takeaways "
+        "regarding {topic}? Provide an overview of {topic}."
+    ),
+    "default": (
+        "What topics and skills are needed to learn {topic}? "
+        "Provide comprehensive material about {topic}."
+    ),
+}
+
+
+class QueryExpander:
+    """Expand short topics/questions into retrieval-optimized queries."""
+
+    def __init__(
+        self,
+        domain: str = "general",
+        keywords: Optional[List[str]] = None,
+        domain_hint: Optional[str] = None,
+    ) -> None:
+        """
+        Args:
+            domain: Free-form domain label (used in chat expansion context).
+            keywords: Optional domain keywords appended to keyword extraction and
+                used to decide when to add ``domain_hint``.
+            domain_hint: Extra phrase appended when a topic matches a keyword
+                (e.g. "Include relevant regulations and safety procedures.").
+        """
+        self.domain = domain
+        self.keywords = [k.lower() for k in (keywords or [])]
+        self.domain_hint = domain_hint
+
+    def expand_for_generation(
+        self,
+        topic: str,
+        content_type: str = "default",
+        file_names: Optional[List[str]] = None,
+    ) -> str:
+        """Expand a topic into a detailed query for content generation."""
+        prefix = ""
+        if file_names:
+            clean = [_strip_ext(n) for n in file_names[:3]]
+            if clean:
+                prefix = f"Related to documents: {', '.join(clean)}. "
+
+        template = _TEMPLATES.get(content_type, _TEMPLATES["default"])
+        expanded = prefix + template.format(topic=topic)
+
+        if self.domain_hint and self._matches_domain(topic):
+            expanded += f" {self.domain_hint}"
+
+        logger.debug(f"expanded '{topic}' -> '{expanded[:80]}...'")
+        return expanded
+
+    def expand_for_chat(self, question: str, conversation_history: str = "") -> str:
+        """Lightly expand a short chat question for better retrieval."""
+        if len(question.split()) >= 5:
+            return question
+        if conversation_history:
+            return f"{question}. Context: {conversation_history[-200:]}"
+        return f"{question} in {self.domain} context"
+
+    def generate_query_variations(
+        self, query: str, num_variations: int = 3
+    ) -> List[str]:
+        """Produce structural query variations for hybrid search."""
+        variations = [query]
+        lower = query.lower()
+        if "what" not in lower and num_variations >= 2:
+            variations.append(f"What is {query}? Explain {query}.")
+        if "how" not in lower and num_variations >= 3:
+            variations.append(f"How to {query}? Steps for {query}.")
+        if self.keywords and num_variations >= 4:
+            variations.append(f"{query} {' '.join(self.keywords[:5])}")
+        return variations[:num_variations]
+
+    def extract_keywords(self, text: str, top_n: int = 10) -> List[str]:
+        """Naive keyword extraction plus any configured domain keywords present."""
+        words = [w for w in text.lower().split() if w not in _STOP_WORDS and len(w) > 3]
+        domain_kws = [k for k in self.keywords if k in text.lower()]
+        seen: List[str] = []
+        for w in words + domain_kws:
+            if w not in seen:
+                seen.append(w)
+        return seen[:top_n]
+
+    def _matches_domain(self, topic: str) -> bool:
+        if not self.keywords:
+            return False
+        low = topic.lower()
+        return any(k in low for k in self.keywords)
+
+
+def _strip_ext(name: str) -> str:
+    for ext in (".pdf", ".docx", ".pptx", ".xlsx", ".txt", ".doc"):
+        name = name.replace(ext, "")
+    return name
+
+
+@lru_cache(maxsize=16)
+def _cached_expander(
+    domain: str, keywords: Tuple[str, ...], domain_hint: Optional[str]
+) -> QueryExpander:
+    return QueryExpander(domain=domain, keywords=list(keywords), domain_hint=domain_hint)
+
+
+def get_query_expander(
+    domain: str = "general",
+    keywords: Optional[List[str]] = None,
+    domain_hint: Optional[str] = None,
+) -> QueryExpander:
+    """Return a cached query expander, one per (domain, keywords, hint) combo.
+
+    Keying on the arguments avoids the single-global footgun where the first
+    caller's domain/keywords would be returned for every subsequent call.
+    """
+    return _cached_expander(domain, tuple(keywords or ()), domain_hint)

aitoolkit/rag/retriever.py

@@ -0,0 +1,141 @@
+"""RAG retriever with optional Redis caching and pluggable filtering.
+
+Differences from the old implementation:
+
+* No LangChain dependency — operates on :class:`RetrievedChunk`.
+* The hardcoded "drop DOCX files" rule is gone; pass a ``filter_fn`` predicate to
+  apply any project-specific filtering instead.
+* Redis is optional (``aitoolkit[cache]``); without it, caching is a no-op.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+from typing import Callable, List, Optional
+
+from loguru import logger
+
+from aitoolkit.rag.vector_store import UnifiedVectorStore
+from aitoolkit.types import RetrievedChunk
+
+try:  # optional dependency
+    import redis.asyncio as redis
+except ImportError:  # pragma: no cover
+    redis = None  # type: ignore[assignment]
+
+FilterFn = Callable[[RetrievedChunk], bool]
+
+
+class RAGRetriever:
+    """Retrieve relevant chunks, with optional caching, filtering and reranking."""
+
+    def __init__(
+        self,
+        vector_store: UnifiedVectorStore,
+        redis_url: Optional[str] = None,
+        cache_ttl: int = 3600,
+        filter_fn: Optional[FilterFn] = None,
+        reranker: Optional[Callable[[str, List[RetrievedChunk]], List[RetrievedChunk]]] = None,
+    ) -> None:
+        self.vector_store = vector_store
+        self.cache_ttl = cache_ttl
+        self.filter_fn = filter_fn
+        self.reranker = reranker
+
+        self.redis_client = None
+        if redis_url and redis is not None:
+            try:
+                self.redis_client = redis.from_url(redis_url, decode_responses=True)
+                logger.info("RAG retriever caching enabled")
+            except Exception as exc:  # noqa: BLE001
+                logger.warning(f"failed to init Redis cache: {exc}")
+        elif redis_url and redis is None:
+            logger.warning(
+                "redis_url provided but redis not installed; install aitoolkit[cache]"
+            )
+
+    def _cache_key(self, query: str, file_ids: Optional[List[str]], limit: int) -> str:
+        payload = {
+            "q": query,
+            "f": sorted(file_ids) if file_ids else None,
+            "l": limit,
+        }
+        digest = hashlib.sha256(json.dumps(payload, sort_keys=True).encode()).hexdigest()
+        return f"aitoolkit:rag:{digest}"
+
+    async def _cache_get(self, key: str) -> Optional[List[RetrievedChunk]]:
+        if not self.redis_client:
+            return None
+        try:
+            raw = await self.redis_client.get(key)
+            if raw:
+                return [RetrievedChunk(**item) for item in json.loads(raw)]
+        except Exception as exc:  # noqa: BLE001
+            logger.warning(f"cache read failed: {exc}")
+        return None
+
+    async def _cache_set(self, key: str, chunks: List[RetrievedChunk]) -> None:
+        if not self.redis_client:
+            return
+        try:
+            raw = json.dumps([c.model_dump() for c in chunks])
+            await self.redis_client.setex(key, self.cache_ttl, raw)
+        except Exception as exc:  # noqa: BLE001
+            logger.warning(f"cache write failed: {exc}")
+
+    async def retrieve(
+        self,
+        query: str,
+        *,
+        file_ids: Optional[List[str]] = None,
+        limit: int = 10,
+        score_threshold: Optional[float] = None,
+        use_cache: bool = True,
+    ) -> List[RetrievedChunk]:
+        """Return relevant chunks for ``query``."""
+        if use_cache:
+            cached = await self._cache_get(self._cache_key(query, file_ids, limit))
+            if cached is not None:
+                logger.debug(f"cache hit ({len(cached)} chunks)")
+                return cached
+
+        chunks = await self.vector_store.similarity_search(
+            query, file_ids=file_ids, limit=limit, score_threshold=score_threshold
+        )
+
+        if self.filter_fn:
+            chunks = [c for c in chunks if self.filter_fn(c)]
+        if self.reranker and chunks:
+            chunks = self.reranker(query, chunks)
+
+        if use_cache and chunks:
+            await self._cache_set(self._cache_key(query, file_ids, limit), chunks)
+        return chunks
+
+    async def get_context_text(
+        self,
+        query: str,
+        *,
+        file_ids: Optional[List[str]] = None,
+        limit: int = 10,
+        separator: str = "\n\n---\n\n",
+        include_sources: bool = True,
+    ) -> str:
+        """Retrieve and format chunks into a single context string."""
+        chunks = await self.retrieve(query, file_ids=file_ids, limit=limit)
+        if not chunks:
+            return ""
+        parts = []
+        for c in chunks:
+            if include_sources:
+                parts.append(
+                    f"[Source: {c.file_id or 'unknown'}, Score: {c.score:.3f}]\n{c.text}"
+                )
+            else:
+                parts.append(c.text)
+        return separator.join(parts)
+
+    async def aclose(self) -> None:
+        if self.redis_client:
+            await self.redis_client.aclose()

aitoolkit/rag/vector_store.py

@@ -0,0 +1,245 @@
+"""Unified Qdrant vector store with file-based filtering.
+
+LangChain-free: it works with plain ``texts`` + ``metadatas`` and returns
+:class:`~aitoolkit.types.RetrievedChunk` objects. Embeddings are produced by an
+injected :class:`~aitoolkit.embeddings.EmbeddingsClient`, so the vector size is
+detected at runtime rather than hardcoded.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+from uuid import uuid4
+
+from loguru import logger
+from qdrant_client import AsyncQdrantClient
+from qdrant_client.http import models as qmodels
+
+from aitoolkit.config import get_settings
+from aitoolkit.embeddings import EmbeddingsClient, get_embeddings_client
+from aitoolkit.exceptions import VectorStoreError
+from aitoolkit.types import RetrievedChunk
+
+
+class UnifiedVectorStore:
+    """Async Qdrant wrapper for storing and retrieving document embeddings."""
+
+    def __init__(
+        self,
+        qdrant_url: Optional[str] = None,
+        collection_name: Optional[str] = None,
+        embeddings: Optional[EmbeddingsClient] = None,
+        vector_size: Optional[int] = None,
+    ) -> None:
+        settings = get_settings()
+        self.qdrant_url = qdrant_url or settings.qdrant_url
+        self.collection_name = collection_name or settings.qdrant_collection
+        self.embeddings = embeddings or get_embeddings_client()
+        self._client = AsyncQdrantClient(
+            url=self.qdrant_url,
+            check_compatibility=settings.qdrant_check_compatibility,
+        )
+        self._initialized = False
+        self._vector_size: Optional[int] = vector_size or settings.qdrant_vector_size
+        logger.info(
+            f"UnifiedVectorStore ready (collection='{self.collection_name}', "
+            f"url={self.qdrant_url})"
+        )
+
+    async def _ensure_collection(self, vector_size: Optional[int] = None) -> None:
+        """Create the collection if it does not exist (no destructive recreate)."""
+        if self._initialized:
+            return
+
+        try:
+            exists = await self._client.collection_exists(self.collection_name)
+        except Exception as exc:  # noqa: BLE001
+            raise VectorStoreError(f"failed to check collection: {exc}") from exc
+
+        if exists:
+            self._initialized = True
+            try:
+                info = await self._client.get_collection(self.collection_name)
+                size = info.config.params.vectors.size  # type: ignore[union-attr]
+                if isinstance(size, int):
+                    self._vector_size = size
+            except Exception:  # noqa: BLE001 - size discovery is best-effort
+                pass
+            return
+
+        size = vector_size or self._vector_size
+        if size is None:
+            # Nothing to create yet — caller will provide a size once embeddings exist.
+            return
+
+        try:
+            await self._client.create_collection(
+                collection_name=self.collection_name,
+                vectors_config=qmodels.VectorParams(
+                    size=size, distance=qmodels.Distance.COSINE
+                ),
+            )
+            self._initialized = True
+            self._vector_size = size
+            logger.success(
+                f"Created collection '{self.collection_name}' (size={size})"
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise VectorStoreError(f"failed to create collection: {exc}") from exc
+
+    async def add_texts(
+        self,
+        texts: List[str],
+        metadatas: Optional[List[Dict[str, Any]]] = None,
+        *,
+        file_id: str,
+        source_type: str = "upload",
+        **extra_metadata: Any,
+    ) -> List[str]:
+        """Embed and store ``texts`` with associated metadata. Returns point IDs."""
+        if not texts:
+            return []
+
+        metadatas = metadatas or [{} for _ in texts]
+        if len(metadatas) != len(texts):
+            raise VectorStoreError("texts and metadatas length mismatch")
+
+        embeddings = await self.embeddings.aembed_documents(texts)
+        if not embeddings:
+            logger.warning("No embeddings produced; nothing stored")
+            return []
+
+        await self._ensure_collection(len(embeddings[0]))
+
+        point_ids: List[str] = []
+        points: List[qmodels.PointStruct] = []
+        for text, vector, meta in zip(texts, embeddings, metadatas):
+            pid = str(uuid4())
+            point_ids.append(pid)
+            points.append(
+                qmodels.PointStruct(
+                    id=pid,
+                    vector=vector,
+                    payload={
+                        "text": text,
+                        "file_id": file_id,
+                        "source_type": source_type,
+                        **meta,
+                        **extra_metadata,
+                    },
+                )
+            )
+
+        try:
+            await self._client.upsert(
+                collection_name=self.collection_name, points=points
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise VectorStoreError(f"failed to upsert points: {exc}") from exc
+
+        logger.info(f"Stored {len(points)} chunks for file_id={file_id}")
+        return point_ids
+
+    async def delete_by_file_id(self, file_id: str) -> None:
+        """Delete all points for a given ``file_id``."""
+        await self._ensure_collection()
+        if not self._initialized:
+            return
+        try:
+            await self._client.delete(
+                collection_name=self.collection_name,
+                points_selector=qmodels.FilterSelector(
+                    filter=qmodels.Filter(
+                        must=[
+                            qmodels.FieldCondition(
+                                key="file_id",
+                                match=qmodels.MatchValue(value=file_id),
+                            )
+                        ]
+                    )
+                ),
+            )
+            logger.info(f"Deleted points for file_id={file_id}")
+        except Exception as exc:  # noqa: BLE001
+            raise VectorStoreError(f"failed to delete by file_id: {exc}") from exc
+
+    async def similarity_search(
+        self,
+        query: str,
+        *,
+        file_ids: Optional[List[str]] = None,
+        limit: int = 10,
+        score_threshold: Optional[float] = None,
+    ) -> List[RetrievedChunk]:
+        """Vector similarity search with optional ``file_ids`` filtering."""
+        query_vector = await self.embeddings.aembed_query(query)
+        await self._ensure_collection(len(query_vector))
+        if not self._initialized:
+            return []
+
+        qfilter = None
+        if file_ids:
+            qfilter = qmodels.Filter(
+                must=[
+                    qmodels.FieldCondition(
+                        key="file_id", match=qmodels.MatchAny(any=file_ids)
+                    )
+                ]
+            )
+
+        try:
+            response = await self._client.query_points(
+                collection_name=self.collection_name,
+                query=query_vector,
+                limit=limit,
+                query_filter=qfilter,
+                score_threshold=score_threshold,
+            )
+        except Exception as exc:  # noqa: BLE001
+            raise VectorStoreError(f"similarity search failed: {exc}") from exc
+
+        chunks: List[RetrievedChunk] = []
+        for point in response.points:
+            payload = dict(point.payload or {})
+            text = payload.pop("text", "")
+            chunks.append(
+                RetrievedChunk(
+                    text=text,
+                    score=point.score,
+                    file_id=payload.get("file_id"),
+                    metadata=payload,
+                )
+            )
+        logger.debug(f"similarity_search returned {len(chunks)} chunks")
+        return chunks
+
+    async def get_unique_file_ids(self, exclude: Optional[List[str]] = None) -> List[str]:
+        """Return all distinct ``file_id`` values stored in the collection."""
+        await self._ensure_collection()
+        if not self._initialized:
+            return []
+        exclude_set = set(exclude or [])
+        found: set[str] = set()
+        offset = None
+        try:
+            while True:
+                records, offset = await self._client.scroll(
+                    collection_name=self.collection_name,
+                    limit=1000,
+                    offset=offset,
+                    with_payload=True,
+                    with_vectors=False,
+                )
+                for rec in records:
+                    fid = (rec.payload or {}).get("file_id")
+                    if fid and fid not in exclude_set:
+                        found.add(fid)
+                if offset is None:
+                    break
+        except Exception as exc:  # noqa: BLE001
+            logger.error(f"failed to scroll file_ids: {exc}")
+            return []
+        return sorted(found)
+
+    async def aclose(self) -> None:
+        await self._client.close()