add rag_tools

rag_tools.py

@@ -1,41 +1,18 @@
 from langchain.tools import tool
-from qdrant_client import QdrantClient
+from vector_store import QdrantStore
-from qdrant_client.http import models
+from chunker import get_chunks
-from langchain_ollama import OllamaEmbeddings
-from langchain_text_splitter import RecursiveCharacterTextSplitter
 
-# Initialize global store
+store = QdrantStore()
-client = QdrantClient(url="http://localhost:6333")
-collection_name = "rag_collection"
-# Ensure collection exists
-if collection_name not in client.get_collections().collections:
-    client.recreate_collection(
-        collection_name=collection_name,
-        vectors_config=models.VectorParams(size=384, distance=models.Distance.COSINE),
-    )
 
-embeddings = OllamaEmbeddings(model="nomic-embed-text")
+@tool
-splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+def search_knowledge_base(query: str, max_results: int = 5) -> str:
+    """Semantic search in knowledge base."""
+    results = store.search(query, limit=max_results)
+    return "\n".join([f"{i+1}. {r['metadata']['title']} – {r['metadata']['content'][:200]}..." for i,r in enumerate(results)])
 
-@tool("search_knowledge_base")
+@tool
-def search_knowledge_base(query: str, max_results: int = 5):
+def add_to_knowledge_base(content: str, title: str) -> str:
-    """Semantic search in the knowledge base."""
+    """Add document to knowledge base."""
-    query_vec = embeddings.embed_query(query)
+    chunks = get_chunks(content, title)
-    results = client.search(
+    store.add_documents(chunks)
-        collection_name=collection_name,
+    return f"Added {len(chunks)} chunks for '{title}'."
-        query_vector=query_vec,
-        limit=max_results,
-        with_payload=True,
-    )
-    return [r.payload for r in results]
-
-@tool("add_to_knowledge_base")
-def add_to_knowledge_base(content: str, title: str):
-    """Add a document to the knowledge base."""
-    chunks = splitter.split_text(content)
-    vectors = embeddings.embed_documents(chunks)
-    points = []
-    for i, (chunk, vec) in enumerate(zip(chunks, vectors)):
-        points.append(models.PointStruct(id=i, vector=vec, payload={"title": title, "chunk": chunk}))
-    client.upsert(collection_name=collection_name, points=models.Batch(points=points))
-    return f"Added {len(chunks)} chunks to the knowledge base."