Add rag_tools

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