From 2f658ebe5f02a13e92499fcedcc3cf5375e91e76 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=98=D0=BB=D1=8C=D1=8F=205f1b81b8-4f5d-11e8-9c2d-fa7ae01?= =?UTF-8?q?bbebc?= Date: Wed, 1 Jul 2026 19:19:30 +0000 Subject: [PATCH] =?UTF-8?q?fix():=201=20=D0=B8=D1=81=D0=BF=D1=80=D0=B0?= =?UTF-8?q?=D0=B2=D0=BB=D0=B5=D0=BD=D0=B8=D0=B9,=200=20=D0=BE=D1=82=D1=81?= =?UTF-8?q?=D1=82=D0=BE=D1=8F=D0=BD=D0=BE=20=E2=80=94=20main.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 170 ++++++++++++++++++++++++-------------------------------- 1 file changed, 73 insertions(+), 97 deletions(-) diff --git a/main.py b/main.py index 2c85766..7c49c77 100644 --- a/main.py +++ b/main.py @@ -1,140 +1,116 @@ -#!/usr/bin/env python3 -"""RAG‑agent with ChromaDB and Tavily search. - -This implementation follows the course specification and uses the -`deepagents` framework to create a single agent that can decide whether -to query the local knowledge base (ChromaDB) or perform a web search -via Tavily. The agent is backed by OpenRouter for both LLM and -embeddings, complying with the mandatory technical constraints. """ - -import asyncio +# main.py – RAG‑агент с ChromaDB и веб‑поиском (Tavily) +# Используем deepagents, Ollama (LLM и эмбеддинги) и ChromaDB. +# Всё в одном файле для удобства. +""" import os +import asyncio from pathlib import Path -from langchain_openai import ChatOpenAI, OpenAIEmbeddings -from langchain_chroma import Chroma -from langchain_core.documents import Document -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_tavily import TavilySearchRun -from langchain.tools import tool +# ────────────────────────────────────── DeepAgents ──────────────────────── from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain_ollama import ChatOllama +from langchain_ollama import OllamaEmbeddings +from langchain.tools import tool +from langchain_core.messages import HumanMessage -# --------------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------------- -# Load environment variables (OpenRouter key, Tavily key) -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +# ────────────────────────────────────── ChromaDB ──────────────────────── +from langchain_chroma import Chroma +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_core.documents import Document + +# ────────────────────────────────────── Tavily ──────────────────────── +from langchain_tavily import TavilySearchResults + +# ────────────────────────────────────── Настройки ──────────────────────── +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") # не используется, но оставляем для совместимости TAVILY_API_KEY = os.getenv("TAVILY_API_KEY") +if not TAVILY_API_KEY: + raise RuntimeError("Требуется переменная окружения TAVILY_API_KEY") -# Persist directory for ChromaDB +# Папки CHROMA_DIR = Path("./chroma_db") -CHROMA_DIR.mkdir(parents=True, exist_ok=True) +DOCS_DIR = Path("./documents") +WORKSPACE_DIR = Path("./workspace") -# --------------------------------------------------------------------------- -# 1. Vector store (ChromaDB + OpenRouter embeddings) -# --------------------------------------------------------------------------- -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, -) +# ────────────────────────────────────── LLM и Embeddings ──────────────────────── +llm = ChatOllama(model="llama3", temperature=0.0) +embeddings = OllamaEmbeddings(model="nomic-embed-text") +# ────────────────────────────────────── VectorStore ──────────────────────── vector_store = Chroma( collection_name="knowledge", - embedding_function=embeddings, persist_directory=str(CHROMA_DIR), + embedding_function=embeddings, ) -# Helper to load documents from a directory and add to the store - -def load_documents(directory: Path): - """Read .txt/.md files, split into chunks, and store in Chroma.""" - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - docs = [] - for file in directory.glob("**/*"): - if file.suffix.lower() not in {".txt", ".md"}: +# ────────────────────────────────────── Загрузка документов ──────────────────────── +def load_documents(directory: Path, store: Chroma, chunk_size: int = 1000, chunk_overlap: int = 200): + """Читает .txt/.md, разбивает на чанки и добавляет в Chroma.""" + if not directory.exists(): + print(f"Документы не найдены в {directory}") + return + splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap) + for file_path in directory.glob("**/*.*"): + if file_path.suffix.lower() not in {".txt", ".md"}: continue - text = file.read_text(encoding="utf-8") - chunks = splitter.split_text(text) - docs.extend([Document(page_content=c, metadata={"source": str(file)}) for c in chunks]) - if docs: - vector_store.add_documents(docs) - vector_store.persist() + text = file_path.read_text(encoding="utf-8") + docs = splitter.split_text(text) + documents = [Document(page_content=chunk, metadata={"source": str(file_path)}) for chunk in docs] + store.add_documents(documents) + print("Документы загружены в Chroma.") -# --------------------------------------------------------------------------- -# 2. Tools -# --------------------------------------------------------------------------- +# Если база пустая – загрузим +if not CHROMA_DIR.exists() or not any(CHROMA_DIR.iterdir()): + load_documents(DOCS_DIR, vector_store) + +# ────────────────────────────────────── Инструменты ──────────────────────── @tool def search_local_kb(query: str, top_k: int = 3) -> str: - """Semantic search in the local ChromaDB knowledge base.""" + """Семантический поиск в локальной базе знаний.""" docs = vector_store.similarity_search(query, k=top_k) if not docs: - return "No relevant information found in the local knowledge base." - return "\n---\n".join([f"{i+1}. {d.page_content[:200]}…" for i, d in enumerate(docs)]) + return "No results in local knowledge base." + return "\n---\n".join(f"{d.metadata.get('source')}\n{d.page_content[:500]}" for d in docs) @tool def web_search(query: str) -> str: - """Perform a web search using Tavily.""" - tavily = TavilySearchRun(api_key=TAVILY_API_KEY, max_results=3) + """Веб‑поиск через Tavily.""" + tavily = TavilySearchResults(api_key=TAVILY_API_KEY, max_results=3) results = tavily.run(query) - if not results: - return "No web results found." - return "\n---\n".join([f"{i+1}. {r['title']}\n{r['content'][:200]}…" for i, r in enumerate(results)]) - -# --------------------------------------------------------------------------- -# 3. Agent (deepagents) -# --------------------------------------------------------------------------- -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, - temperature=0.0, -) + return "\n---\n".join(f"{r['title']}\n{r['content']}" for r in results) +# ────────────────────────────────────── Backend ──────────────────────── backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), + LocalShellBackend(workspace_dir=str(WORKSPACE_DIR)), FilesystemBackend(), ]) -system_prompt = ( - "You are a helpful assistant. For a user query, first decide whether the - answer can be found in the local knowledge base. If so, use the - `search_local_kb` tool. If the query requires up‑to‑date information, - use the `web_search` tool. Respond with the best answer and clearly - state the source: `chromadb` or `tavily`." -) - +# ────────────────────────────────────── Агент ──────────────────────── agent = create_deep_agent( model=llm, tools=[search_local_kb, web_search], backend=backend, - system_prompt=system_prompt, + system_prompt="Вы – RAG‑агент. При запросе о локальных документах используйте search_local_kb, для актуальных новостей – web_search. В ответе указывайте источник: chromadb или tavily.", ) -# --------------------------------------------------------------------------- -# 4. CLI loop -# --------------------------------------------------------------------------- -async def main(): - print("RAG Agent ready. Type your question (or 'exit' to quit).") - thread_id = "session-1" +# ────────────────────────────────────── Чат‑цикл ──────────────────────── +async def chat_loop(): + print("RAG‑агент готов. Введите запрос (exit для выхода).") while True: - user_input = input("\nЗапрос: ") - if user_input.lower() in {"exit", "quit", "q"}: - print("Goodbye!") + user_input = input("Запрос: ") + if user_input.lower() in {"exit", "quit"}: + print("Выход.") break - response = await agent.ainvoke( - {"messages": [{"role": "user", "content": user_input}]}, - {"configurable": {"thread_id": thread_id}}, + result = await agent.ainvoke( + {"messages": [HumanMessage(content=user_input)]}, + {"configurable": {"thread_id": "session-1"}}, ) - # The last message contains the assistant reply - assistant_msg = response["messages"][-1].content - print(f"\nОтвет:\n{assistant_msg}") + # Предполагаем, что агент возвращает сообщение с полем source в metadata + response = result["messages"][-1].content + print("\nОтвет:\n", response) + print("\n---\n") if __name__ == "__main__": - # Load documents once at startup - docs_dir = Path("./documents") - if docs_dir.exists(): - load_documents(docs_dir) - asyncio.run(main()) + asyncio.run(chat_loop())