commit 6195155d67c9afa7141bdc36640d1012d4ff6fe1 Author: Аделина Саттарова Date: Thu May 28 13:23:49 2026 +0000 Add rag_tools diff --git a/rag_tools.py b/rag_tools.py new file mode 100644 index 0000000..a4d456e --- /dev/null +++ b/rag_tools.py @@ -0,0 +1,31 @@ +from langchain.tools import tool +from qdrant_client import QdrantClient +from langchain.embeddings.ollama import OllamaEmbeddings +from langchain.vectorstores.qdrant import QdrantVectorStore +from langchain.text_splitter import RecursiveCharacterTextSplitter + +# Initialize embeddings and vector store +embeddings = OllamaEmbeddings(model="nomic-embed-text") +client = QdrantClient(host="localhost", port=6333) +collection_name = "knowledge_base" +# Ensure collection exists +if not client.has_collection(collection_name): + client.create_collection(name=collection_name, vectors_config={"size": embeddings.embed_query(["test"]).shape[1], "distance": "Cosine"}) +vector_store = QdrantVectorStore(client=client, collection_name=collection_name, embedding=embeddings) + +# Text splitter +splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + +@tool("search_knowledge_base") +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """Search the knowledge base for relevant documents.""" + results = vector_store.similarity_search_with_score(query, k=max_results) + return "\n---\n".join([f"{score:.4f}: {doc.page_content[:200]}..." for doc, score in results]) + +@tool("add_to_knowledge_base") +def add_to_knowledge_base(content: str, title: str) -> str: + """Add a new document to the knowledge base.""" + chunks = splitter.split_text(content) + docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] + vector_store.add_documents(docs) + return f"Added {len(chunks)} chunks from '{title}'."