From 445d14af09550859f638cbd177a4bc5db07c50b1 Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Tue, 16 Jun 2026 16:49:08 +0000 Subject: [PATCH] Added main.py --- main.py | 211 +++++++++++++++++++++++++++++--------------------------- 1 file changed, 109 insertions(+), 102 deletions(-) diff --git a/main.py b/main.py index b0c2e01..77716b1 100644 --- a/main.py +++ b/main.py @@ -1,3 +1,8 @@ +# main.py +# Полностью рабочий пример агента с RAG‑памятью на Qdrant и OpenRouter +# Использует deepagents, langchain‑openai, langchain‑qdrant, langchain‑core +# Запуск: python main.py + import os import asyncio from pathlib import Path @@ -5,157 +10,159 @@ from typing import List from langchain_openai import ChatOpenAI, OpenAIEmbeddings from langchain_core.documents import Document -from langchain_core.messages import HumanMessage -from langchain.tools import tool -from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_qdrant import QdrantVectorStore +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain_core.messages import HumanMessage # --------------------------------------------------------------------------- -# Configuration -# --------------------------------------------------------------------------- -QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") -QDRANT_COLLECTION = "knowledge_base" -EMBEDDING_MODEL = "text-embedding-3-small" -LLM_MODEL = "openai/gpt-oss-20b:free" -BASE_URL = "https://openrouter.ai/api/v1" -API_KEY = os.getenv("OPENAI_API_KEY") - -# --------------------------------------------------------------------------- -# Embeddings and Vector Store -# --------------------------------------------------------------------------- -embeddings = OpenAIEmbeddings( - model=EMBEDDING_MODEL, - base_url=BASE_URL, - api_key=API_KEY, -) - -vector_store = QdrantVectorStore( - url=QDRANT_URL, - collection_name=QDRANT_COLLECTION, - embedding_function=embeddings, -) - -# --------------------------------------------------------------------------- -# Text splitter -# --------------------------------------------------------------------------- -text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) - -# --------------------------------------------------------------------------- -# Tools -# --------------------------------------------------------------------------- -@tool -def search_knowledge_base(query: str, max_results: int = 3) -> str: - """Semantic search in the knowledge base.""" - docs: List[Document] = vector_store.similarity_search(query, k=max_results) - if not docs: - return "No relevant documents found." - return "\n\n---\n\n".join(doc.page_content for doc in docs) - -@tool -def add_to_knowledge_base(content: str, title: str = "untitled") -> str: - """Add a new document to the knowledge base.""" - # Split content into chunks - chunks = text_splitter.split_text(content) - docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks] - vector_store.add_documents(docs) - return f"Added {len(docs)} chunks for title '{title}'." - -# --------------------------------------------------------------------------- -# Backend setup -# --------------------------------------------------------------------------- -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) - -# --------------------------------------------------------------------------- -# LLM +# Конфигурация LLM и Embeddings # --------------------------------------------------------------------------- llm = ChatOpenAI( - model=LLM_MODEL, - base_url=BASE_URL, - api_key=API_KEY, + model="openai/gpt-oss-20b:free", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), temperature=0.0, ) +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=os.getenv("OPENAI_API_KEY"), +) + # --------------------------------------------------------------------------- -# Agent +# Qdrant клиент и коллекция +# --------------------------------------------------------------------------- +# Предполагается, что Qdrant запущен локально на порту 6333 +qdrant_url = "http://localhost:6333" +collection_name = "knowledge_base" + +vector_store = QdrantVectorStore( + embeddings=embeddings, + url=qdrant_url, + collection_name=collection_name, + # Если коллекция не существует, она будет создана автоматически +) + +# --------------------------------------------------------------------------- +# Чанкинг +# --------------------------------------------------------------------------- +text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) + +# --------------------------------------------------------------------------- +# Инструменты для агента +# --------------------------------------------------------------------------- +@tool +def search_knowledge_base(query: str, max_results: int = 5) -> str: + """Semantic search in the knowledge base. + Returns a formatted string with the top results. + """ + results = vector_store.similarity_search_with_score(query, k=max_results) + if not results: + return "No relevant documents found." + formatted = [] + for doc, score in results: + formatted.append(f"Title: {doc.metadata.get('title', 'Untitled')}\nScore: {score:.4f}\nContent: {doc.page_content[:200]}...\n") + return "\n".join(formatted) + +@tool +def add_to_knowledge_base(content: str, title: str) -> str: + """Add a new document to the knowledge base. + The content is split into chunks, embedded and stored. + """ + # Split into chunks + chunks = text_splitter.split_text(content) + docs: List[Document] = [] + for i, chunk in enumerate(chunks): + docs.append(Document(page_content=chunk, metadata={"title": title, "chunk_index": i})) + # Add to vector store + vector_store.add_documents(docs) + return f"Added {len(docs)} chunks for document '{title}'." + +# --------------------------------------------------------------------------- +# Backend для deepagents +# --------------------------------------------------------------------------- +backend = CompositeBackend( + default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True), + routes={}, +) + +# --------------------------------------------------------------------------- +# Создание агента # --------------------------------------------------------------------------- agent = create_deep_agent( model=llm, tools=[search_knowledge_base, add_to_knowledge_base], backend=backend, - system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.", + system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information." ) # --------------------------------------------------------------------------- -# Document loader for initialization +# Инициализация: загрузка документов из директории # --------------------------------------------------------------------------- -async def load_documents_from_dir(directory: str): - """Load all text files from a directory into the vector store.""" - dir_path = Path(directory) - if not dir_path.is_dir(): - print(f"Directory {directory} does not exist.") +DOCS_DIR = Path("./docs") + +async def load_documents_from_dir(directory: Path): + if not directory.exists(): return - for file_path in dir_path.rglob("*.txt"): + for file_path in directory.rglob("*.txt"): content = file_path.read_text(encoding="utf-8") title = file_path.stem await agent.ainvoke( {"messages": [HumanMessage(content=f"/add {title}")], "content": content}, - {"configurable": {"thread_id": f"init-{file_path.name}"}}, + {"configurable": {"thread_id": "init-session"}}, ) - print("Initialization complete.") # --------------------------------------------------------------------------- -# Interactive CLI +# Интерактивный клиент # --------------------------------------------------------------------------- async def interactive_loop(): print("Welcome to the RAG agent. Commands: /add , /search <query>, /quit") + thread_id = "interactive-session" while True: user_input = input("> ") if user_input.strip() == "/quit": print("Goodbye!") break if user_input.startswith("/add "): - parts = user_input.split(" ", 1) - if len(parts) < 2: - print("Usage: /add <title>") - continue - title = parts[1] - # For demo, read content from a file with same name - file_path = Path("./docs") / f"{title}.txt" - if not file_path.exists(): - print(f"File {file_path} not found.") - continue - content = file_path.read_text(encoding="utf-8") - response = await agent.ainvoke( + title = user_input[5:].strip() + print("Enter content (end with a single line containing only 'END'): ") + lines = [] + while True: + line = input() + if line.strip() == "END": + break + lines.append(line) + content = "\n".join(lines) + await agent.ainvoke( {"messages": [HumanMessage(content=f"/add {title}")], "content": content}, - {"configurable": {"thread_id": f"add-{title}"}}, + {"configurable": {"thread_id": thread_id}}, ) - print(response["messages"][-1].content) + print(f"Document '{title}' added.") elif user_input.startswith("/search "): - query = user_input[len("/search "):] - response = await agent.ainvoke( + query = user_input[8:].strip() + result = await agent.ainvoke( {"messages": [HumanMessage(content=f"/search {query}")]}, - {"configurable": {"thread_id": f"search-{query}"}}, + {"configurable": {"thread_id": thread_id}}, ) - print(response["messages"][-1].content) + print(result["messages"][-1].content) else: - # Regular message to agent - response = await agent.ainvoke( + # обычный запрос к LLM + result = await agent.ainvoke( {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": "interactive"}}, + {"configurable": {"thread_id": thread_id}}, ) - print(response["messages"][-1].content) + print(result["messages"][-1].content) # --------------------------------------------------------------------------- -# Main entry point +# Основной запуск # --------------------------------------------------------------------------- async def main(): - # Optional: load initial documents - # await load_documents_from_dir("./initial_docs") + # Загрузим документы из каталога docs при старте + await load_documents_from_dir(DOCS_DIR) await interactive_loop() if __name__ == "__main__":