42 lines
1.1 KiB
Python
42 lines
1.1 KiB
Python
from langchain.tools import tool
|
|
from .vector_store import VectorStore
|
|
from .splitter import chunk_text
|
|
|
|
# Shared store instance
|
|
_store: VectorStore | None = None
|
|
|
|
@tool
|
|
def add_to_knowledge_base(content: str, title: str = "Document") -> str:
|
|
"""Add content to the knowledge base.
|
|
|
|
Parameters
|
|
----------
|
|
content: str
|
|
Text content to add.
|
|
title: str
|
|
Optional title for the document.
|
|
"""
|
|
global _store
|
|
if _store is None:
|
|
_store = VectorStore()
|
|
chunks = chunk_text(content)
|
|
_store.add_documents(chunks)
|
|
return f"Added {len(chunks)} chunks to the knowledge base under title '{title}'."
|
|
|
|
@tool
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> list[tuple[str, float]]:
|
|
"""Search the knowledge base for relevant chunks.
|
|
|
|
Parameters
|
|
----------
|
|
query: str
|
|
Search query.
|
|
max_results: int
|
|
Number of top results to return.
|
|
"""
|
|
global _store
|
|
if _store is None:
|
|
_store = VectorStore()
|
|
results = _store.search(query, max_results=max_results)
|
|
return results
|