From c9bae1e62713785eed180cc91e8f4325a7133aeb 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:45:35 +0000 Subject: [PATCH] fix: main.py --- main.py | 86 +++++++++++++++++++++++++++------------------------------ 1 file changed, 41 insertions(+), 45 deletions(-) diff --git a/main.py b/main.py index f990f5b..f560db4 100644 --- a/main.py +++ b/main.py @@ -1,18 +1,21 @@ 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 -# Load environment variables +# --------------------- 1. Загрузка переменных окружения --------------------- load_dotenv() -# ---------- LLM and Embeddings ---------- +# --------------------- 2. LLM и Embeddings --------------------- llm = ChatOpenAI( model="openai/gpt-oss-20b:free", base_url="https://openrouter.ai/api/v1", @@ -26,85 +29,78 @@ embeddings = OpenAIEmbeddings( api_key=os.getenv("OPENAI_API_KEY"), ) -# ---------- Vector Store ---------- -CHROMA_DIR = "./chroma_db" +# --------------------- 3. Векторное хранилище Chroma --------------------- +CHROMA_DIR = Path("./chroma_db") +CHROMA_DIR.mkdir(parents=True, exist_ok=True) vector_store = Chroma( collection_name="knowledge", embedding_function=embeddings, - persist_directory=CHROMA_DIR, + persist_directory=str(CHROMA_DIR), ) -# ---------- Document Loader ---------- - -def load_documents(directory: str): - """Read .txt/.md files, split into chunks and add to Chroma.""" +# --------------------- 4. Загрузка документов --------------------- +DOCS_DIR = Path("./documents") +if DOCS_DIR.exists(): splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - docs = [] - for root, _, files in os.walk(directory): - for file in files: - if file.lower().endswith(('.txt', '.md')): - path = os.path.join(root, file) - with open(path, 'r', encoding='utf-8') as f: - text = f.read() - chunks = splitter.split_text(text) - docs.extend([Document(page_content=c, metadata={"source": path}) for c in chunks]) - if docs: - vector_store.add_documents(docs) - vector_store.persist() + 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() -# Load documents once at startup -if not os.path.exists(CHROMA_DIR) or not os.listdir(CHROMA_DIR): - load_documents("./documents") - -# ---------- Tools ---------- +# --------------------- 5. Инструменты --------------------- @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 information found in local knowledge base." - return "\n---\n".join(f"{d.metadata.get('source', 'unknown')}\n{d.page_content}" for d in docs) + return "No relevant local knowledge found." + return "\n---\n".join(f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)) @tool def web_search(query: str) -> str: """Web search using Tavily.""" - from langchain_tavily import TavilySearchResults - tavily = TavilySearchResults(max_results=3) + tavily = TavilySearchResults(api_key=os.getenv("TAVILY_API_KEY")) results = tavily.run(query) if not results: return "No web results found." - return "\n---\n".join(f"{r['title']}\n{r['content']}" for r in results) + return "\n---\n".join(f"{i+1}. {res['title']}\n{res['url']}\n{res['content'][:200]}..." for i, res in enumerate(results)) -# ---------- Backend ---------- +# --------------------- 6. Backend --------------------- backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) -# ---------- Agent ---------- +# --------------------- 7. Создание агента --------------------- agent = create_deep_agent( model=llm, tools=[search_local_kb, web_search], backend=backend, - system_prompt="You are a helpful assistant. For questions about local documents use the local knowledge base. For up‑to‑date facts use web search. Always state the source (chromadb or tavily) in your answer.", + 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`." + ), ) -# ---------- Main Loop ---------- -async def main(): - print("RAG Agent ready. Type 'exit' to quit.") +# --------------------- 8. CLI --------------------- +async def chat_loop(): + print("Добро пожаловать в RAG‑агент. Введите 'exit' для выхода.") while True: user_input = input("\nЗапрос: ") - if user_input.lower() in {"exit", "quit"}: - print("Goodbye!") + if user_input.lower() in {"exit", "quit", "q"}: + print("До свидания!") break - # Invoke agent result = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, + {"messages": [{"role": "user", "content": user_input}]}, {"configurable": {"thread_id": "session-1"}}, ) - # Extract last message content - answer = result["messages"][-1].content - print(f"\nОтвет:\n{answer}") + # Последнее сообщение агента + content = result["messages"][-1]["content"] + print(content) if __name__ == "__main__": - asyncio.run(main()) + asyncio.run(chat_loop())