From 0491a6b732a0d1a3a8af6901bad41d86c82fbf5b Mon Sep 17 00:00:00 2001 From: furyhawk Date: Mon, 15 Jun 2026 10:02:56 +0800 Subject: [PATCH] feat: implement RAG search tools and add related documentation --- backend/core/agent.py | 13 +++ skills/__init__.py | 0 skills/rag_search/SKILL.md | 31 +++++++ skills/rag_search/__init__.py | 0 skills/rag_search/search.py | 166 ++++++++++++++++++++++++++++++++++ 5 files changed, 210 insertions(+) create mode 100644 skills/__init__.py create mode 100644 skills/rag_search/SKILL.md create mode 100644 skills/rag_search/__init__.py create mode 100644 skills/rag_search/search.py diff --git a/backend/core/agent.py b/backend/core/agent.py index ab53d12..fc028c4 100644 --- a/backend/core/agent.py +++ b/backend/core/agent.py @@ -10,6 +10,7 @@ from pydantic_ai import Agent from pydantic_ai.capabilities import MCP, Thinking, ToolSearch, WebSearch from pydantic_ai.models.openai import OpenAIResponsesModel from pydantic_ai.providers.openai import OpenAIProvider +from pydantic_ai.tools import Tool from pydantic_ai_harness import CodeMode # Community packages, alphabetical: @@ -62,6 +63,7 @@ class AgentService: self._agent = Agent( self._model, capabilities=self._build_capabilities(), + tools=self._build_rag_tools(), ) async def shutdown(self) -> None: @@ -89,6 +91,17 @@ class AgentService: # ── Internal helpers ─────────────────────────────────────────────────── + @staticmethod + def _build_rag_tools() -> list[Tool]: + """Build RAG search tools for document retrieval.""" + from skills.rag_search.search import rag_search, rag_search_by_document, rag_list_collections + + return [ + Tool(rag_search, name="rag_search"), + Tool(rag_search_by_document, name="rag_search_by_document"), + Tool(rag_list_collections, name="rag_list_collections"), + ] + def _build_capabilities(self) -> list: """Assemble the full list of agent capabilities.""" return [ diff --git a/skills/__init__.py b/skills/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/skills/rag_search/SKILL.md b/skills/rag_search/SKILL.md new file mode 100644 index 0000000..c67b817 --- /dev/null +++ b/skills/rag_search/SKILL.md @@ -0,0 +1,31 @@ +--- +name: rag-search +description: "Search ingested documents using vector retrieval. Use this when the user asks about information that may be in uploaded documents (PDFs, DOCX, TXT, MD files)." +tags: [rag, retrieval, documents, search] +version: "1.0.0" +--- + +# RAG Search + +Retrieves relevant document chunks from the vector database using semantic search. +Results include the chunk content, relevance score, source filename, and page number. + +## When to use + +- The user asks a question about documents they've uploaded +- The user wants to find information across ingested files +- The user references "my documents", "the PDF", "the uploaded file", etc. + +## How it works + +1. The query is embedded using the configured embedding model +2. Milvus performs a cosine similarity search against all document chunks +3. Top results are returned with content, score, and source metadata + +## Example + +``` +User: "What does the report say about revenue growth?" +Agent: [calls rag_search(query="revenue growth") → gets relevant chunks] +Agent: "Based on the documents, the revenue grew by 15% in Q3..." +``` diff --git a/skills/rag_search/__init__.py b/skills/rag_search/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/skills/rag_search/search.py b/skills/rag_search/search.py new file mode 100644 index 0000000..7427def --- /dev/null +++ b/skills/rag_search/search.py @@ -0,0 +1,166 @@ +"""RAG document search tools for the agent. + +Provides async tool functions that search ingested documents via Milvus +vector store and return relevant chunks with source metadata. +""" + +from __future__ import annotations + +from typing import Any + +from backend.core.config import settings +from backend.services.rag.embeddings import EmbeddingService +from backend.services.rag.retrieval import RetrievalService +from backend.services.rag.vectorstore import MilvusVectorStore as VectorStore + +# ── Lazy singleton ────────────────────────────────────────────────────────── + +_retrieval: RetrievalService | None = None +_vector_store: VectorStore | None = None + + +def _get_services() -> tuple[RetrievalService | None, VectorStore | None]: + """Build RetrievalService + VectorStore singletons (lazy).""" + global _retrieval, _vector_store + if _retrieval is not None: + return _retrieval, _vector_store + try: + rag_settings = settings.rag + embed_service = EmbeddingService(settings=rag_settings) + _vector_store = VectorStore(settings=rag_settings, embedding_service=embed_service) + _retrieval = RetrievalService(vector_store=_vector_store, settings=rag_settings) + return _retrieval, _vector_store + except Exception as exc: + import logging + + logging.getLogger(__name__).warning("RAG retrieval unavailable: %s", exc) + return None, None + + +# ── Tool functions ────────────────────────────────────────────────────────── + + +async def rag_search( + query: str, + collection_name: str = "documents", + limit: int = 5, +) -> list[dict[str, Any]]: + """Search ingested documents for information relevant to the query. + + Use this tool when the user asks about content in uploaded documents + (PDFs, DOCX, TXT, MD files). Returns a list of relevant text chunks + with their source filename, page number, and relevance score. + + Args: + query: The search question or keywords to look up in the documents. + collection_name: Which document collection to search (default: "documents"). + limit: Maximum number of results to return (1-10, default: 5). + + Returns: + A list of result dicts, each with keys: + - content: The matching text chunk + - score: Relevance score (0.0-1.0, higher is better) + - filename: Source document filename + - filetype: File type (pdf, docx, txt, md) + - page_num: Page number in the source document (if applicable) + - chunk_num: Chunk sequence number within the document + """ + retrieval, _ = _get_services() + if retrieval is None: + return [ + { + "content": "", + "score": 0.0, + "error": "RAG search is not available — check EMBEDDING_BASE_URL configuration.", + } + ] + + results = await retrieval.retrieve( + query=query, + collection_name=collection_name, + limit=min(limit, 10), + min_score=0.0, + ) + + return [ + { + "content": r.content, + "score": round(r.score, 4), + "filename": r.metadata.get("filename", "unknown"), + "filetype": r.metadata.get("filetype", ""), + "page_num": r.metadata.get("page_num", 0), + "chunk_num": r.metadata.get("chunk_num", 0), + "source_path": r.metadata.get("source_path", ""), + } + for r in results + ] + + +async def rag_search_by_document( + query: str, + filename: str, + collection_name: str = "documents", + limit: int = 3, +) -> list[dict[str, Any]]: + """Search within a specific uploaded document by filename. + + Use this when the user wants to find information in a particular file. + + Args: + query: The search question or keywords. + filename: The filename to restrict results to (e.g., "report.pdf"). + collection_name: Which document collection to search (default: "documents"). + limit: Maximum number of results to return (1-10, default: 3). + + Returns: + A list of result dicts with content, score, filename, and source metadata. + """ + retrieval, _ = _get_services() + if retrieval is None: + return [ + { + "content": "", + "score": 0.0, + "error": "RAG search is not available — check EMBEDDING_BASE_URL configuration.", + } + ] + + results = await retrieval.retrieve( + query=query, + collection_name=collection_name, + limit=min(limit * 3, 30), # fetch more for post-filtering + min_score=0.0, + ) + + # Filter results by filename + filtered = [ + r for r in results if r.metadata.get("filename", "").lower() == filename.lower() + ] + + return [ + { + "content": r.content, + "score": round(r.score, 4), + "filename": r.metadata.get("filename", "unknown"), + "filetype": r.metadata.get("filetype", ""), + "page_num": r.metadata.get("page_num", 0), + "chunk_num": r.metadata.get("chunk_num", 0), + } + for r in filtered[:limit] + ] + + +async def rag_list_collections() -> list[str]: + """List available document collections that can be searched. + + Returns: + A list of collection names (e.g., ["documents"]). + """ + try: + rag_settings = settings.rag + embed_service = EmbeddingService(settings=rag_settings) + store = VectorStore(settings=rag_settings, embedding_service=embed_service) + collections = await store.list_collections() + return collections + except Exception: + return []