29 lines
1.2 KiB
Python
29 lines
1.2 KiB
Python
from langchain.tools import tool
|
|
from langchain_chroma import Chroma
|
|
from langchain_ollama import OllamaEmbeddings
|
|
from langchain_tavily import TavilySearchResults
|
|
import os
|
|
from dotenv import load_dotenv
|
|
|
|
# Загружаем переменные окружения из .env
|
|
load_dotenv()
|
|
|
|
@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."""
|
|
tavily_api_key = os.getenv("TAVILY_API_KEY")
|
|
if not tavily_api_key:
|
|
raise ValueError("Требуется переменная окружения TAVILY_API_KEY")
|
|
search = TavilySearchResults(tavily_api_key=tavily_api_key)
|
|
results = search.run(query)
|
|
return "\n\n".join([result.get("content", "") for result in results]) |