From 51fe3adb3e892e19d3c70423cd4d0b1bb38719cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=93=D0=BB=D0=B5=D0=B1=20=D0=9D=D0=B8=D0=BA=D0=B8=D1=88?= =?UTF-8?q?=D0=B8=D0=BD?= Date: Thu, 4 Jun 2026 16:57:53 +0000 Subject: [PATCH] fix: main.py --- main.py | 75 ++++++++++++++++++++++++++++----------------------------- 1 file changed, 37 insertions(+), 38 deletions(-) diff --git a/main.py b/main.py index f560db4..ee6117e 100644 --- a/main.py +++ b/main.py @@ -2,20 +2,18 @@ import os import asyncio from pathlib import Path from dotenv import load_dotenv - from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_chroma import Chroma from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_core.documents import Document -from langchain_tavily import TavilySearchResults from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend -# --------------------- 1. Загрузка переменных окружения --------------------- +# Load environment variables load_dotenv() -# --------------------- 2. LLM и Embeddings --------------------- +# ---------- LLM and Embeddings ---------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -29,78 +27,79 @@ embeddings = OpenAIEmbeddings( api_key=os.getenv("OPENAI_API_KEY"), ) -# --------------------- 3. Векторное хранилище Chroma --------------------- +# ---------- Vector Store ---------- CHROMA_DIR = Path("./chroma_db") -CHROMA_DIR.mkdir(parents=True, exist_ok=True) vector_store = Chroma( collection_name="knowledge", embedding_function=embeddings, persist_directory=str(CHROMA_DIR), ) -# --------------------- 4. Загрузка документов --------------------- -DOCS_DIR = Path("./documents") -if DOCS_DIR.exists(): - splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - for file_path in DOCS_DIR.rglob("*.txt"): - 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] - vector_store.add_documents(documents) - vector_store.persist() +# ---------- Document Loader ---------- -# --------------------- 5. Инструменты --------------------- +def load_documents(directory: str, vectorstore: Chroma): + splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + docs = [] + for file_path in Path(directory).glob("**/*.*"): + if file_path.suffix.lower() not in {".txt", ".md"}: + continue + text = file_path.read_text(encoding="utf-8") + chunks = splitter.split_text(text) + docs.extend([Document(page_content=c, metadata={"source": str(file_path)}) for c in chunks]) + if docs: + vectorstore.add_documents(docs) + vectorstore.persist() + +# Load initial documents if collection is empty +if not CHROMA_DIR.exists() or not list(CHROMA_DIR.iterdir()): + load_documents("documents", vector_store) + +# ---------- Tools ---------- @tool def search_local_kb(query: str, top_k: int = 3) -> str: """Semantic search in the local knowledge base.""" docs = vector_store.similarity_search(query, k=top_k) if not docs: - return "No relevant local knowledge found." - return "\n---\n".join(f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)) + return "No relevant information found in local knowledge base." + return "\n---\n".join([f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs]) @tool def web_search(query: str) -> str: """Web search using Tavily.""" - tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) + from langchain_tavily import TavilySearchResults + tavily = TavilySearchResults() results = tavily.run(query) - if not results: - return "No web results found." - return "\n---\n".join(f"{i+1}. {res['title']}\n{res['url']}\n{res['content'][:200]}..." for i, res in enumerate(results)) + return "\n---\n".join([f"{r['title']}\n{r['content']}" for r in results]) -# --------------------- 6. Backend --------------------- +# ---------- Backend ---------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# --------------------- 7. Создание агента --------------------- +# ---------- Agent ---------- agent = create_deep_agent( model=llm, tools=[search_local_kb, web_search], backend=backend, - system_prompt=( - "You are an assistant that answers user questions.\n" - "If the question is about information that should be in the local knowledge base, use the tool `search_local_kb`.\n" - "If the question requires up‑to‑date information from the web, use the tool `web_search`.\n" - "Always indicate the source of the answer in the format: `Источник: chromadb` or `Источник: tavily`." - ), + system_prompt="You are a helpful assistant that can search both a local knowledge base and the web.\nWhen answering, always indicate the source: either 'chromadb' or 'tavily'.\nUse the appropriate tool based on the query context.", ) -# --------------------- 8. CLI --------------------- -async def chat_loop(): - print("Добро пожаловать в RAG‑агент. Введите 'exit' для выхода.") +# ---------- CLI ---------- +async def main(): + print("RAG Agent with ChromaDB and Tavily. Type 'exit' to quit.") while True: user_input = input("\nЗапрос: ") - if user_input.lower() in {"exit", "quit", "q"}: - print("До свидания!") + if user_input.lower() in {"exit", "quit"}: + print("Goodbye!") break result = await agent.ainvoke( {"messages": [{"role": "user", "content": user_input}]}, {"configurable": {"thread_id": "session-1"}}, ) - # Последнее сообщение агента + # Extract last message content content = result["messages"][-1]["content"] print(content) if __name__ == "__main__": - asyncio.run(chat_loop()) + asyncio.run(main())