From 041b6358cd138563e80df7c322c64739acccb6a0 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:27:32 +0000 Subject: [PATCH] add tools.py --- tools.py | 43 +++++++++++++++++++------------------------ 1 file changed, 19 insertions(+), 24 deletions(-) diff --git a/tools.py b/tools.py index 47be5fe..64e4c16 100644 --- a/tools.py +++ b/tools.py @@ -1,38 +1,33 @@ """ -Tools for the RAG agent. +RAG tools for the agent. -Two tools: search_knowledge_base and add_to_knowledge_base. +search_knowledge_base and add_to_knowledge_base are decorated with @tool. """ - +import os 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") +@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: + """Semantic search in the vector store.""" + results = vector_store.similarity_search(query, k=max_results) + if not results: 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) + out_lines = [] + for i, res in enumerate(results, 1): + out_lines.append(f"{i}. {res['content'][:200]}... (distance: {res['distance']:.3f})") + return "\n".join(out_lines) -@tool("add_to_knowledge_base") +@tool("Add document 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. - """ + """Adds a text chunk to the vector store.""" + # Split content into chunks chunks = split_text(content) + docs = [] 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}'." + docs.append({"content": chunk, "metadata": {"title": title, "chunk_index": idx}}) + vector_store.add_documents(docs) + return f"Added {len(chunks)} chunks to the knowledge base."