diff --git a/tools.py b/tools.py new file mode 100644 index 0000000..b32b757 --- /dev/null +++ b/tools.py @@ -0,0 +1,22 @@ +from langchain.tools import tool +from langchain_community.vectorstores import Chroma +from langchain_ollama import OllamaEmbeddings +from langchain_tavily import TavilySearchResults +import os + +@tool +def search_local_kb(query: str, top_k: int = 3) -> str: + """Поиск в локальной базе знаний (ChromaDB) по запросу.""" + persist_directory = "./chroma_db" + embeddings = OllamaEmbeddings(model="nomic-embed-text") + vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings) + retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) + docs = retriever.get_relevant_documents(query) + return "\n\n".join([doc.page_content for doc in docs]) + +@tool +def web_search(query: str) -> str: + """Поиск в интернете через Tavily.""" + search = TavilySearchResults() + results = search.run(query) + return "\n\n".join([result.get("content", "") for result in results]) \ No newline at end of file