83 lines
2.4 KiB
Python
83 lines
2.4 KiB
Python
"""Tools for RAG agent.
|
||
|
||
This module defines two LangChain tools:
|
||
|
||
* ``search_knowledge_base`` – semantic search in Qdrant.
|
||
* ``add_to_knowledge_base`` – add a document to Qdrant.
|
||
|
||
The tools use the global ``vector_store`` instance defined in this module.
|
||
"""
|
||
|
||
from typing import List
|
||
|
||
from langchain_ollama import OllamaEmbeddings
|
||
from langchain_qdrant import QdrantVectorStore
|
||
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
||
from langchain.tools import tool
|
||
from langchain.schema import Document
|
||
|
||
# Global vector store instance
|
||
# We initialise it lazily – the first call to the tools will create the store.
|
||
_vector_store = None
|
||
|
||
# Embedding model
|
||
_embedding = OllamaEmbeddings(model="nomic-embed-text")
|
||
|
||
# Text splitter
|
||
_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||
|
||
|
||
def _get_vector_store() -> QdrantVectorStore:
|
||
global _vector_store
|
||
if _vector_store is None:
|
||
# Create a new collection named "knowledge_base"
|
||
_vector_store = QdrantVectorStore.from_texts(
|
||
texts=[],
|
||
embedding=_embedding,
|
||
url="http://localhost:6333",
|
||
collection_name="knowledge_base",
|
||
)
|
||
return _vector_store
|
||
|
||
|
||
@tool("search_knowledge_base")
|
||
def search_knowledge_base(query: str, max_results: int = 5) -> List[dict]:
|
||
"""Semantic search in the knowledge base.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The search query.
|
||
max_results: int
|
||
Maximum number of results to return.
|
||
|
||
Returns
|
||
-------
|
||
List[dict]
|
||
List of matching documents with ``content`` and ``metadata``.
|
||
"""
|
||
store = _get_vector_store()
|
||
results = store.similarity_search_with_score(query, k=max_results)
|
||
# Return only content and metadata for simplicity
|
||
return [
|
||
{"content": doc.page_content, "metadata": doc.metadata}
|
||
for doc, _ in results
|
||
]
|
||
|
||
|
||
@tool("add_to_knowledge_base")
|
||
def add_to_knowledge_base(content: str, title: str) -> str:
|
||
"""Add a document to the knowledge base.
|
||
|
||
The content is split into chunks before being added.
|
||
"""
|
||
store = _get_vector_store()
|
||
chunks = _splitter.split_text(content)
|
||
docs: List[Document] = []
|
||
for i, chunk in enumerate(chunks):
|
||
docs.append(
|
||
Document(page_content=chunk, metadata={"title": title, "chunk_index": i})
|
||
)
|
||
store.add_documents(docs)
|
||
return f"Document '{title}' added with {len(chunks)} chunks."
|