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}'."