diff --git a/main.py b/main.py index f9d2ded..8f8a025 100644 --- a/main.py +++ b/main.py @@ -1,41 +1,143 @@ import os -from dotenv import load_dotenv +from pathlib import Path -from vectorstore import create_vectorstore, load_documents -from agent import create_agent +# Загрузка переменных окружения (TAVILY_API_KEY) +try: + from dotenv import load_dotenv -def main(): - load_dotenv() # загружает TAVILY_API_KEY из .env + load_dotenv() +except Exception: # pragma: no cover + pass - persist_dir = "./chroma_db" - vectorstore = create_vectorstore(persist_directory=persist_dir) +# Основные зависимости LangChain и Ollama +from langchain_ollama import ChatOllama, OllamaEmbeddings +from langchain_chroma import Chroma +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.tools import tool +from langchain.agents import AgentExecutor, create_openai_tools_agent +from langchain_core.messages import HumanMessage - # Загружаем документы только если база пустая - if vectorstore._collection.count() == 0: - docs_dir = "./documents" - load_documents(docs_dir, vectorstore) +# Импорт вспомогательных функций из других файлов проекта +try: + from vectorstore import create_vectorstore, load_documents # noqa: F401 +except Exception: # pragma: no cover + # Если модули не найдены – создаём простые заглушки для запуска примера + def create_vectorstore(persist_directory="./chroma_db"): + return Chroma( + persist_directory=persist_directory, + embedding_function=OllamaEmbeddings(model="nomic-embed-text"), + ) - agent_executor = create_agent(vectorstore) + def load_documents(directory, vectorstore): + pass + +# --------------------------------------------------------------------------- # +# 1. Создаём векторное хранилище и загружаем документы (если ещё нет) +# --------------------------------------------------------------------------- # + +CHROMA_DIR = "./chroma_db" +DOCS_DIR = "./documents" + +vectorstore = create_vectorstore(persist_directory=CHROMA_DIR) + +# Если коллекция пустая – загрузим документы из каталога +if not vectorstore.get_ids(): + load_documents(DOCS_DIR, vectorstore) + +# --------------------------------------------------------------------------- # +# 2. Определяем инструменты агента +# --------------------------------------------------------------------------- # + +@tool("search_local_kb") +def search_local_kb(query: str, top_k: int = 3) -> str: + """ + Семантический поиск по локальной базе знаний (ChromaDB). + Возвращает объединённый текст найденных сегментов. + """ + retriever = vectorstore.as_retriever(search_kwargs={"k": top_k}) + docs = retriever.get_relevant_documents(query) + if not docs: + return "Ничего не найдено в локальной базе." + # Собираем контент из документов + content = "\n\n".join(doc.page_content for doc in docs) + return f"[Local KB]\n{content}\nИсточник: chromadb" + +@tool("web_search") +def web_search(query: str) -> str: + """ + Поиск в интернете через Tavily. + Возвращает краткое резюме найденных результатов. + """ + from tavily import TavilyClient + + client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY")) + results = client.search(query=query, max_results=3) + if not results: + return "Ничего не найдено в интернете." + snippets = "\n\n".join( + f"{res['title']}\n{res['content']}" for res in results + ) + return f"[Web Search]\n{snippets}\nИсточник: tavily" + +# --------------------------------------------------------------------------- # +# 3. Создаём агента с выбором источника +# --------------------------------------------------------------------------- # + +llm = ChatOllama(model="llama3") + +tools = [search_local_kb, web_search] + +agent = create_openai_tools_agent( + llm=llm, + tools=tools, + system_message=( + "Ты помощник, отвечающий на вопросы. " + "Если вопрос относится к содержимому локальных документов, используй инструмент search_local_kb; " + "если требуется актуальная информация из интернета – web_search. " + "В ответе обязательно указывай источник (chromadb или tavily)." + ), +) + +agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=False) + +# --------------------------------------------------------------------------- # +# 4. CLI чат‑цикл +# --------------------------------------------------------------------------- # + +def main() -> None: + print("=== RAG-агент с ChromaDB и Tavily ===") + print("Введите 'exit' для выхода.\n") - print("RAG-агент готов. Введите 'exit' для выхода.") while True: - query = input("\nЗапрос: ").strip() - if query.lower() in ("exit", "quit"): + try: + query = input("Запрос: ").strip() + except (KeyboardInterrupt, EOFError): + print("\nВыход.") break - if not query: - continue - result = agent_executor.invoke({"input": query}) - # Ожидаем, что агент вернёт dict с полями 'output' и, опционально, 'source' - if isinstance(result, dict): - output = result.get("output", "") - source = result.get("source", "unknown") - else: - output = str(result) - source = "unknown" + if not query or query.lower() == "exit": + print("До свидания!") + break - print(f"\n{output}") - print(f"Источник: {source}") + # Запускаем агента + response = agent_executor.invoke({"input": query}) + answer = response.get("output", "") + + # Выводим ответ и источник (если он есть в тексте) + source_line = "" + if "[Local KB]" in answer: + source_line = "Источник: chromadb" + elif "[Web Search]" in answer: + source_line = "Источник: tavily" + + print("\nОтвет:") + # Убираем маркеры из ответа + cleaned_answer = answer.replace("[Local KB]\n", "").replace( + "[Web Search]\n", "" + ) + print(cleaned_answer.strip()) + if source_line: + print(f"\n{source_line}\n") if __name__ == "__main__": main() \ No newline at end of file