add tools
This commit is contained in:
@@ -0,0 +1,20 @@
|
||||
from langchain.tools import tool
|
||||
from langchain.embeddings.ollama import OllamaEmbeddings
|
||||
from langchain.vectorstores import Chroma
|
||||
import os
|
||||
|
||||
@tool("search_local_kb")
|
||||
def search_local_kb(query: str, top_k: int = 3) -> str:
|
||||
"""Semantic search in local ChromaDB."""
|
||||
persist_dir = os.getenv("CHROMA_DB_DIR", "./chroma_db")
|
||||
embeddings = OllamaEmbeddings(model="nomic-embed-text")
|
||||
db = Chroma(embedding_function=embeddings, persist_directory=persist_dir)
|
||||
docs = db.similarity_search(query, k=top_k)
|
||||
return "\n".join([d.page_content for d in docs]) + f"\n[Source: chromadb]"
|
||||
|
||||
@tool("web_search")
|
||||
def web_search(query: str) -> str:
|
||||
from tavily import TavilyClient
|
||||
client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))
|
||||
result = client.search(query, max_results=3)
|
||||
return "\n".join([r['content'] for r in result]) + f"\n[Source: tavily]"
|
||||
Reference in New Issue
Block a user