71 lines
1.6 KiB
Python
71 lines
1.6 KiB
Python
from typing import List
|
|
|
|
from langchain.tools import tool
|
|
from langchain_ollama import Ollama
|
|
from langchain_qdrant import QdrantVectorStore
|
|
from langchain_tavily import TavilySearchResults
|
|
|
|
|
|
@tool
|
|
def search_local_kb(
|
|
query: str,
|
|
top_k: int,
|
|
vectorstore: QdrantVectorStore,
|
|
) -> List[str]:
|
|
"""
|
|
Perform a semantic search in the local knowledge base stored in Qdrant.
|
|
|
|
Parameters
|
|
----------
|
|
query : str
|
|
The user's query.
|
|
top_k : int
|
|
Number of top results to return.
|
|
vectorstore : QdrantVectorStore
|
|
The vector store to search.
|
|
|
|
Returns
|
|
-------
|
|
List[str]
|
|
List of relevant document snippets.
|
|
"""
|
|
retriever = vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.get_relevant_documents(query)
|
|
return [doc.page_content for doc in docs]
|
|
|
|
|
|
@tool
|
|
def web_search(
|
|
query: str,
|
|
tavily_api_key: str,
|
|
max_results: int = 3,
|
|
) -> List[str]:
|
|
"""
|
|
Perform a web search using Tavily.
|
|
|
|
Parameters
|
|
----------
|
|
query : str
|
|
The user's query.
|
|
tavily_api_key : str
|
|
Tavily API key.
|
|
max_results : int
|
|
Number of search results to return.
|
|
|
|
Returns
|
|
-------
|
|
List[str]
|
|
List of search result snippets.
|
|
"""
|
|
tavily = TavilySearchResults(
|
|
api_key=tavily_api_key,
|
|
max_results=max_results,
|
|
)
|
|
results = tavily.run(query)
|
|
# Extract snippets from results
|
|
snippets = []
|
|
for result in results:
|
|
snippet = result.get("content") or result.get("snippet") or result.get("title")
|
|
if snippet:
|
|
snippets.append(snippet)
|
|
return snippets |