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])