From 8298cac7341ba59118e5cbfa8fc9934ca237bc04 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=9A=D1=83=D1=82?= =?UTF-8?q?=D0=BB=D0=B0=D1=85=D0=BC=D0=B5=D1=82=D0=BE=D0=B2?= Date: Thu, 28 May 2026 09:15:09 +0000 Subject: [PATCH] Add tools.py --- tools.py | 38 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 38 insertions(+) create mode 100644 tools.py diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..47be5fe --- /dev/null +++ b/tools.py @@ -0,0 +1,38 @@ +""" +Tools for the RAG agent. + +Two tools: search_knowledge_base and add_to_knowledge_base. +""" + +from typing import List, Dict +from langchain.tools import tool +from vector_store import vector_store +from chunker import split_text + +@tool("search_knowledge_base") +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """Semantic search in the knowledge base. + + Returns a formatted string of results. + """ + hits = vector_store.search(query, k=max_results) + if not hits: + return "No relevant documents found." + lines: List[str] = [] + for i, hit in enumerate(hits, 1): + title = hit["metadata"].get("title", f"doc_{hit['id']}") + snippet = hit["document"][:200] + lines.append(f"{i}. {title}: {snippet}...") + return "\n".join(lines) + +@tool("add_to_knowledge_base") +def add_to_knowledge_base(content: str, title: str = "document") -> str: + """Add a document to the knowledge base. + + Splits content into chunks and stores each with metadata. + """ + chunks = split_text(content) + for idx, chunk in enumerate(chunks): + doc_id = f"{title}_{idx}" + vector_store.add_document(doc_id=doc_id, text=chunk, metadata={"title": title}) + return f"Added {len(chunks)} chunks from '{title}'."