replypier/backend/app/tools/knowledge_search_byagent.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import traceback
from typing import Any, Dict, Iterable, List, Optional, Set, Tuple
from agno.tools import Toolkit
from agno.utils.log import logger
from sqlalchemy import text
from app.database import SessionLocal
from app.core.config import settings
from app.knowledge.chunking import chunk_id, page_id_of, split_text
from app.repositories.knowledge import KnowledgeRepository
from agno.knowledge.agent import AgentKnowledge
from agno.vectordb.pgvector import PgVector, SearchType
from app.knowledge.embedder import get_embedder
from uuid import UUID

NO_RESULTS = "No relevant information found in the knowledge base."
NO_FURTHER_RESULTS = (
    "No further results for this topic in the knowledge base beyond what earlier "
    "searches in this turn already returned."
)


class _Hit:
    """One chunk row as rendered to the model, plus what groups and orders it."""

    __slots__ = ("id", "content", "page_id", "chunk", "chunk_count", "url", "name")

    def __init__(self, id: str, content: str, meta: Dict[str, Any], name: str):
        self.id = id
        self.content = content
        self.name = name
        self.page_id = meta.get("page_id") or page_id_of(id)
        self.chunk = meta.get("chunk") if isinstance(meta.get("chunk"), int) else None
        self.chunk_count = meta.get("chunk_count") if isinstance(meta.get("chunk_count"), int) else None
        self.url = meta.get("url")

    @property
    def is_chunked(self) -> bool:
        return self.chunk is not None and self.chunk_count is not None and self.chunk_count > 1


class KnowledgeSearchByAgent(Toolkit):
    def __init__(self, agent_id: str, org_id: UUID, source: str = None):
        super().__init__(name="knowledge_search_by_agent")
        self.name = "knowledge_search_by_agent"
        self.description = "Search the knowledge base for information about a query"
        self.function = self.search_knowledge_base
        self.agent_id = agent_id
        self.org_id = org_id
        self.source = source
        # Structured citations for the most recent turn: list of {"name", "type"}.
        # Read (and reset) by the chat agent after each run to attach to the response.
        self.collected_sources: List[Dict[str, str]] = []
        # Chunk ids already shown to the model this turn. A repeat search skips
        # them and returns the next-best material instead of the same text again,
        # so every tool call within the turn's budget adds something new.
        self.seen_chunk_ids: Set[str] = set()

        # NOTE: this used to export the org's key as OPENAI_API_KEY for agno's
        # old default OpenAI embedder. Search now uses the local FastEmbed
        # embedder and every model gets its key explicitly, so the env write
        # was dead — and a cross-tenant race (concurrent orgs overwrote each
        # other's key in process-global state).
        self.agent_knowledge = None
        self._table: Optional[str] = None
        self._schema: Optional[str] = None
        self.register(self.search_knowledge_base)

    def reset_turn(self) -> None:
        """Forget this turn's citations and shown chunks before the next run."""
        self.collected_sources = []
        self.seen_chunk_ids = set()

    def search_knowledge_base(self, query: str) -> str:
        """Use this function to search the knowledge base for information about a query.

        Args:
            query: The query to search for.
        """
        try:
            logger.debug(f"Searching knowledge base for query: {query}")

            # Use context manager for database operations
            with SessionLocal() as db:
                knowledge_repo = KnowledgeRepository(db)
                # Get knowledge sources linked to this agent
                knowledge_sources = knowledge_repo.get_by_agent(self.agent_id)

                if not knowledge_sources:
                    return "No knowledge sources available for this agent."

                # All of an agent's sources live in the org's one vector table.
                source = knowledge_sources[0]
                self._table = source.table_name
                self._schema = source.schema

                # Initialize agent_knowledge if it doesn't exist
                if self.agent_knowledge is None:
                    # Initialize vector db with simpler search type to avoid connection issues
                    vector_db = PgVector(
                        table_name=source.table_name,
                        db_url=settings.DATABASE_URL,
                        schema=source.schema,
                        search_type=SearchType.vector,  # Changed from hybrid to vector for speed
                        embedder=get_embedder()
                    )
                    logger.debug(f"Vector db initialized: {source.table_name}")

                    # Create AgentKnowledge instance
                    self.agent_knowledge = AgentKnowledge(vector_db=vector_db)

                # Convert UUID to string in filters
                filters = {"agent_id": [str(self.agent_id)]}
                if self.source:
                    filters["name"] = self.source
                logger.debug(f"Search filters: {filters}")

                # Ask for enough rows that, after skipping what this turn has
                # already shown, a full set of new ones remains.
                limit = settings.KNOWLEDGE_SEARCH_RESULTS
                documents = self.agent_knowledge.search(
                    query=query,
                    num_documents=limit + len(self.seen_chunk_ids),
                    filters=filters
                )
                logger.debug(f"Documents: {documents}")

                hits = [
                    _Hit(
                        getattr(doc, "id", None) or f"{doc.name}#{position}",
                        doc.content,
                        getattr(doc, "meta_data", None) or {},
                        doc.name or "Untitled",
                    )
                    for position, doc in enumerate(documents)
                    if doc.content
                ]
                fresh = [hit for hit in hits if hit.id not in self.seen_chunk_ids][:limit]
                if not fresh:
                    return NO_FURTHER_RESULTS if hits else NO_RESULTS

                hits_by_id = {hit.id: hit for hit in fresh}
                neighbours, heads = self._page_context(db, fresh)
                for hit in neighbours:
                    hits_by_id.setdefault(hit.id, hit)

                source_types = {
                    source.source: source.source_type.value.lower() for source in knowledge_sources
                }
                rendered, shown = self._render(hits_by_id.values(), heads, source_types)
                # Only what the model actually received counts as seen or as a
                # citation; a page dropped by the size cap stays searchable.
                self.seen_chunk_ids.update(hit.id for hit in shown)
                self._record_sources(shown, source_types)
                return rendered

        except Exception as e:
            logger.error(f"Error searching knowledge base: {str(e)}")
            logger.error(f"Full traceback: {traceback.format_exc()}")
            return "Error searching knowledge base."

    def _page_context(self, db, hits: List[_Hit]) -> Tuple[List[_Hit], Dict[str, _Hit]]:
        """Context for each hit on a chunked page, fetched in one query.

        Neighbours — the chunk before and after a hit — so a fact split by a
        chunk boundary (a pricing table, a numbered procedure) arrives whole.

        Heads — the page's first chunk, keyed by page id — which is where a
        page says what it is: a product's name and price, a plan's headline
        terms. A query that matches the description further down (or is
        phrased in another language than the page) would otherwise come back
        without them. Heads are optional: ``_render`` adds them only where
        they fit after every matched page has its place.

        Only chunked pages have context; a single-row page is complete."""
        neighbour_ids: Set[str] = set()
        head_ids: Set[str] = set()
        for hit in hits:
            if not hit.is_chunked:
                continue
            for index in (hit.chunk - 1, hit.chunk + 1):
                if 1 <= index <= hit.chunk_count:
                    neighbour_ids.add(chunk_id(hit.page_id, index))
            head_ids.add(chunk_id(hit.page_id, 1))
        known = self.seen_chunk_ids | {hit.id for hit in hits}
        neighbour_ids -= known
        head_ids -= known | neighbour_ids
        wanted = neighbour_ids | head_ids
        if not wanted:
            return [], {}

        rows = db.execute(
            text(
                f'SELECT id, name, content, meta_data FROM {self._schema}."{self._table}" '
                "WHERE id = ANY(:ids)"
            ),
            {"ids": list(wanted)},
        ).fetchall()
        fetched = [_Hit(row.id, row.content, row.meta_data or {}, row.name or "Untitled") for row in rows]
        neighbours = [hit for hit in fetched if hit.id in neighbour_ids]
        heads = {hit.page_id: hit for hit in fetched if hit.id in head_ids}
        return neighbours, heads

    def _record_sources(self, hits: Iterable[_Hit], source_types: Dict[str, str]) -> None:
        # Record structured citations (deduped by name+type) so the chat agent
        # can surface them to the widget.
        seen = {(s['name'], s['type']) for s in self.collected_sources}
        for hit in hits:
            key = (hit.name, source_types.get(hit.name, 'unknown'))
            if key not in seen:
                seen.add(key)
                self.collected_sources.append({'name': key[0], 'type': key[1]})

    def _render(
        self, hits: Iterable[_Hit], heads: Dict[str, _Hit], source_types: Dict[str, str]
    ) -> Tuple[str, List[_Hit]]:
        """Group chunks by page, in page order, under one header per page.

        Pages with matches are placed first and dropped whole once the result
        would exceed the size cap — a page is never cut mid-chunk, and the cap
        keeps one search from crowding out the rest of the conversation. Page
        heads go in afterwards, only into pages that made it and only while
        they still fit, so a head never displaces a page that matched.
        Returns the text and the hits it contains."""
        cap = settings.KNOWLEDGE_SEARCH_MAX_CHARS
        pages: Dict[str, List[_Hit]] = {}
        for hit in hits:
            pages.setdefault(hit.page_id, []).append(hit)

        placed: List[List[_Hit]] = []
        total = 0
        for page_hits in pages.values():
            size = len(_page_block(page_hits, source_types)) + len(_PAGE_SEPARATOR)
            if placed and total + size > cap:
                break
            placed.append(page_hits)
            total += size

        for page_hits in placed:
            head = heads.get(page_hits[0].page_id)
            if head is None:
                continue
            size = len(_chunk_text(head)) + len(_CHUNK_SEPARATOR)
            if total + size <= cap:
                page_hits.append(head)
                total += size

        blocks = [_page_block(page_hits, source_types) for page_hits in placed]
        shown = [hit for page_hits in placed for hit in page_hits]
        return _PAGE_SEPARATOR.join(blocks), shown


_PAGE_SEPARATOR = "\n\n"
_CHUNK_SEPARATOR = "\n"


def _chunk_text(hit: _Hit) -> str:
    label = f"(part {hit.chunk}/{hit.chunk_count}) " if hit.is_chunked else ""
    return f"{label}{_hit_text(hit)}"


def _page_block(page_hits: List[_Hit], source_types: Dict[str, str]) -> str:
    """One page as the model sees it: a source header, then its chunks in
    page order."""
    page_hits.sort(key=lambda h: h.chunk or 0)
    first = page_hits[0]
    header = f"[{source_types.get(first.name, 'unknown').upper()} - {first.name}"
    if first.url and first.url != first.name:
        header += f" | {first.url}"
    header += "]"
    return header + " " + _CHUNK_SEPARATOR.join(_chunk_text(hit) for hit in page_hits)


# A row this much bigger than a chunk was stored before pages were chunked and
# has not been re-indexed yet (scripts/rechunk_knowledge.py). It is a data
# defect, not knowledge to protect: one such row can be 100k+ chars.
OVERSIZED_ROW_FACTOR = 4
OVERSIZED_ROW_NOTE = "\n(page continues — not yet re-indexed in full)"


def _hit_text(hit: _Hit) -> str:
    limit = settings.KNOWLEDGE_CHUNK_SIZE * OVERSIZED_ROW_FACTOR
    if len(hit.content) <= limit:
        return hit.content
    kept: List[str] = []
    total = 0
    for piece in split_text(hit.content, settings.KNOWLEDGE_CHUNK_SIZE):
        if total + len(piece) > limit:
            break
        kept.append(piece)
        total += len(piece)
    return "".join(kept) + OVERSIZED_ROW_NOTE