40 lines
1.4 KiB
Python
40 lines
1.4 KiB
Python
from langchain.tools import tool
|
|
from vector_store import search_documents, add_documents
|
|
|
|
|
|
@tool
|
|
def search_knowledge_base(query: str, max_results: int = 5) -> str:
|
|
"""Semantic search in the knowledge base. Use this tool to find relevant information.
|
|
|
|
Args:
|
|
query: The search query string.
|
|
max_results: Maximum number of results to return (default 5).
|
|
|
|
Returns:
|
|
Formatted string with search results and relevance scores.
|
|
"""
|
|
results = search_documents(query, max_results=max_results)
|
|
if not results:
|
|
return "No results found in the knowledge base."
|
|
output_lines = [f"Found {len(results)} result(s):\n"]
|
|
for i, r in enumerate(results, 1):
|
|
title = r["metadata"].get("title", "Unknown")
|
|
content = r["content"]
|
|
output_lines.append(f"[{i}] Title: {title}")
|
|
output_lines.append(f" {content}\n")
|
|
return "\n".join(output_lines)
|
|
|
|
|
|
@tool
|
|
def add_to_knowledge_base(content: str, title: str) -> str:
|
|
"""Add a document to the knowledge base. Use this tool to store new information.
|
|
|
|
Args:
|
|
content: The full text content of the document to add.
|
|
title: A descriptive title for the document.
|
|
|
|
Returns:
|
|
Confirmation message with the number of chunks stored.
|
|
"""
|
|
num_chunks = add_documents(content=content, title=title)
|
|
return f"Successfully added document '{title}' to the knowledge base ({num_chunks} chunk(s) stored)." |