diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..4a61cb3 --- /dev/null +++ b/tools.py @@ -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]"