Add tools.py
This commit is contained in:
@@ -0,0 +1,37 @@
|
|||||||
|
"""
|
||||||
|
Tools for the RAG agent.
|
||||||
|
|
||||||
|
- search_local_kb(query, top_k)
|
||||||
|
- web_search(query)
|
||||||
|
"""
|
||||||
|
import os
|
||||||
|
from typing import List
|
||||||
|
|
||||||
|
from langchain_ollama import OllamaEmbeddings
|
||||||
|
from langchain_chroma import Chroma
|
||||||
|
from langchain_tavily import TavilySearchResults
|
||||||
|
from langchain_core.documents import Document
|
||||||
|
|
||||||
|
# Global vectorstore – will be set in init_db or main
|
||||||
|
vectorstore: Chroma = None
|
||||||
|
|
||||||
|
def search_local_kb(query: str, top_k: int = 3) -> List[Document]:
|
||||||
|
"""Semantic search in the local ChromaDB."""
|
||||||
|
if vectorstore is None:
|
||||||
|
raise RuntimeError("Vectorstore not initialized")
|
||||||
|
return vectorstore.similarity_search_with_score(query, k=top_k)
|
||||||
|
|
||||||
|
# Tavily client – API key from env
|
||||||
|
from tavily import TavilyClient
|
||||||
|
import os
|
||||||
|
TAVILY_API_KEY = os.getenv("TAVILY_API_KEY", "")
|
||||||
|
client = TavilyClient(api_key=TAVILY_API_KEY)
|
||||||
|
|
||||||
|
def web_search(query: str, max_results: int = 3) -> List[Document]:
|
||||||
|
"""Search the web via Tavily and return Documents."""
|
||||||
|
results = client.search(query=query, max_results=max_results)
|
||||||
|
docs = []
|
||||||
|
for r in results:
|
||||||
|
content = f"{r.title}\n\n{r.content}"
|
||||||
|
docs.append(Document(page_content=content, metadata={"source": "tavily", "url": r.url}))
|
||||||
|
return docs
|
||||||
Reference in New Issue
Block a user