From b9cd98addffc674e28684970a014734989e3c79c 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: Wed, 27 May 2026 14:37:31 +0000 Subject: [PATCH] add tools.py --- tools.py | 71 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 71 insertions(+) create mode 100644 tools.py diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..9861279 --- /dev/null +++ b/tools.py @@ -0,0 +1,71 @@ +""" +LangChain tools for adding and searching a Qdrant knowledge base. + +The module exposes two @tool‑decorated functions: +* add_to_knowledge_base(content: str, title: str = "document") +* search_knowledge_base(query: str, max_results: int = 5) + +Both tools use an in‑memory Qdrant client and Ollama embeddings. +""" +import os +from typing import List + +from langchain_core.documents import Document +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_ollama import ChatOllama, OllamaEmbeddings +from langchain_qdrant import QdrantVectorStore +from langchain.tools import tool + +# --- Configuration ------------------------------------------------------- +QDRANT_COLLECTION = "rag_agent" +EMBEDDING_MODEL = "nomic-embed-text" +LLM_MODEL = "llama3" + +# Initialize embeddings and vector store (in‑memory Qdrant) +embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL, base_url="http://localhost:11434") +vector_store = QdrantVectorStore(client=None, collection_name=QDRANT_COLLECTION, embedding=embeddings) + +# Ensure the collection exists (creates if not present) +try: + vector_store.client.create_collection( + name=QDRANT_COLLECTION, + vectors_config=dict(size=embeddings.embed_query("test").shape[0], distance="cosine"), + ) +except Exception: + # Collection already exists – ignore + pass + +# Text splitter for chunking documents +splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + +@tool("Add a document to the knowledge base") +def add_to_knowledge_base(content: str, title: str = "document") -> str: + """ + Splits *content* into chunks, embeds them via Ollama and stores in Qdrant. + Returns a confirmation string with number of chunks added. + """ + # Split content + docs: List[Document] = splitter.split_text(content) + for i, chunk in enumerate(docs): + doc = Document(page_content=chunk, metadata={"title": title, "chunk_index": i}) + vector_store.add_documents([doc]) + return f"Added {len(docs)} chunks from '{title}'." + +@tool("Search the knowledge base") +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """ + Performs a semantic similarity search in Qdrant. + Returns a formatted string with top results and their metadata. + """ + docs = vector_store.similarity_search(query, k=max_results) + if not docs: + return "No relevant documents found." + lines: List[str] = [] + for i, doc in enumerate(docs, 1): + title = doc.metadata.get("title", "unknown") + idx = doc.metadata.get("chunk_index", "?") + lines.append(f"{i}. [{title} - chunk {idx}]\n{doc.page_content[:200]}...") + return "\n\n".join(lines) + +# Expose tool names for agent creation +TOOLS = [add_to_knowledge_base, search_knowledge_base]