From 771faa84e060a7f4679803e51db76b060f23971a Mon Sep 17 00:00:00 2001 From: Danil Parunin 5f1b81b8-4f5d-11e8-9c2d-fa7ae01bbebc Date: Tue, 16 Jun 2026 16:29:35 +0000 Subject: [PATCH] =?UTF-8?q?fix(needs=5Ffixes):=201=20=D0=B8=D1=81=D0=BF?= =?UTF-8?q?=D1=80=D0=B0=D0=B2=D0=BB=D0=B5=D0=BD=D0=B8=D0=B9,=200=20=D0=BE?= =?UTF-8?q?=D1=82=D1=81=D1=82=D0=BE=D1=8F=D0=BD=D0=BE=20=E2=80=94=20main.p?= =?UTF-8?q?y?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 219 ++++++++++++++++++++++++++++---------------------------- 1 file changed, 109 insertions(+), 110 deletions(-) diff --git a/main.py b/main.py index c135e3a..b0c2e01 100644 --- a/main.py +++ b/main.py @@ -1,163 +1,162 @@ -"""RAG Agent with Qdrant and OpenRouter. - -This script implements the assignment requirements: -* Two tools – `search_knowledge_base` and `add_to_knowledge_base` – are defined with the `@tool` decorator. -* A Qdrant vector store is used for semantic search. Documents are split into chunks with a - `RecursiveCharacterTextSplitter` that has `chunk_overlap=100` as requested. -* The agent is created with LangChain’s `create_agent` (the "Исправить" instruction overrides the - earlier requirement to use `create_deep_agent`). -* A simple CLI allows adding documents, searching the knowledge base and quitting. - -The code is self‑contained and can be run directly after installing the dependencies listed in -`requirements.txt`. -""" - import os import asyncio -import pathlib +from pathlib import Path 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_community.vectorstores import Qdrant -from langchain_community.document_loaders import TextLoader from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain.agents import create_agent, AgentExecutor, AgentType +from langchain_qdrant import QdrantVectorStore +from deepagents import create_deep_agent +from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend # --------------------------------------------------------------------------- # Configuration # --------------------------------------------------------------------------- -OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") -if not OPENAI_API_KEY: - raise RuntimeError("OPENAI_API_KEY environment variable is required") - -# LLM and embeddings via OpenRouter -llm = ChatOpenAI( - model="openai/gpt-oss-20b:free", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, - temperature=0.0, -) +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="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=OPENAI_API_KEY, + model=EMBEDDING_MODEL, + base_url=BASE_URL, + api_key=API_KEY, ) -# --------------------------------------------------------------------------- -# Vector store setup (Qdrant) -# --------------------------------------------------------------------------- -# Qdrant is expected to be running locally on the default port 6333. -# If you need a different host/port, adjust the `url` parameter. -vector_store = Qdrant.from_existing_index( - collection_name="knowledge", - embeddings=embeddings, - url="http://localhost:6333", +vector_store = QdrantVectorStore( + url=QDRANT_URL, + collection_name=QDRANT_COLLECTION, + embedding_function=embeddings, ) # --------------------------------------------------------------------------- # Text splitter # --------------------------------------------------------------------------- -text_splitter = RecursiveCharacterTextSplitter( - chunk_size=1000, - chunk_overlap=100, # as required by the assignment -) +text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) # --------------------------------------------------------------------------- # Tools # --------------------------------------------------------------------------- @tool def search_knowledge_base(query: str, max_results: int = 3) -> str: - """Search the knowledge base for relevant information. - - Parameters - ---------- - query: str - The search query. - max_results: int, optional - Number of top results to return. Defaults to 3. - """ - docs = vector_store.similarity_search(query, k=max_results) + """Semantic search in the knowledge base.""" + docs: List[Document] = vector_store.similarity_search(query, k=max_results) if not docs: - return "No results found." + 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 = "document") -> str: - """Add content to the knowledge base. - - Parameters - ---------- - content: str - The raw text to add. - title: str, optional - A title for the document. Defaults to "document". - """ - # Split into chunks and create Document objects +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) - from langchain_core.documents import Document - 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}'." # --------------------------------------------------------------------------- -# Agent creation (LangChain create_agent) +# Backend setup # --------------------------------------------------------------------------- -# The system prompt instructs the agent to use the knowledge base tools. -SYSTEM_PROMPT = ( - "You are an AI assistant with access to a knowledge base. " - "Use the tools `search_knowledge_base` and `add_to_knowledge_base` to answer user queries. " - "If the user asks to add information, store it. If the user asks for information, search the base." +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) + +# --------------------------------------------------------------------------- +# LLM +# --------------------------------------------------------------------------- +llm = ChatOpenAI( + model=LLM_MODEL, + base_url=BASE_URL, + api_key=API_KEY, + temperature=0.0, ) -agent = create_agent( - llm=llm, +# --------------------------------------------------------------------------- +# Agent +# --------------------------------------------------------------------------- +agent = create_deep_agent( + model=llm, tools=[search_knowledge_base, add_to_knowledge_base], - system_prompt=SYSTEM_PROMPT, - agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, + backend=backend, + system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.", ) -agent_executor = AgentExecutor(agent=agent, tools=[search_knowledge_base, add_to_knowledge_base]) +# --------------------------------------------------------------------------- +# 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.") + return + for file_path in dir_path.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}"}}, + ) + print("Initialization complete.") # --------------------------------------------------------------------------- -# CLI +# Interactive CLI # --------------------------------------------------------------------------- -async def handle_user_input(user_input: str) -> str: - if user_input.startswith("/add "): - # Expected format: /add | <content> - try: - _, rest = user_input.split("/add ", 1) - title, content = rest.split("|", 1) - title = title.strip() - content = content.strip() - result = add_to_knowledge_base(content, title) - return result - except ValueError: - return "Invalid format. Use: /add <title> | <content>" - elif user_input.startswith("/search "): - query = user_input[len("/search "):].strip() - return search_knowledge_base(query) - elif user_input == "/quit": - return "quit" - else: - # Forward to the agent - response = await agent_executor.ainvoke({"messages": [HumanMessage(content=user_input)]}) - return response["messages"][-1].content - -async def main(): - print("RAG Agent CLI. Commands: /add <title> | <content>, /search <query>, /quit") +async def interactive_loop(): + print("Welcome to the RAG agent. Commands: /add <title>, /search <query>, /quit") while True: user_input = input("> ") - if not user_input: - continue - result = await handle_user_input(user_input) - if result == "quit": + if user_input.strip() == "/quit": print("Goodbye!") break - print(result) + 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( + {"messages": [HumanMessage(content=f"/add {title}")], "content": content}, + {"configurable": {"thread_id": f"add-{title}"}}, + ) + print(response["messages"][-1].content) + elif user_input.startswith("/search "): + query = user_input[len("/search "):] + response = await agent.ainvoke( + {"messages": [HumanMessage(content=f"/search {query}")]}, + {"configurable": {"thread_id": f"search-{query}"}}, + ) + print(response["messages"][-1].content) + else: + # Regular message to agent + response = await agent.ainvoke( + {"messages": [HumanMessage(content=user_input)]}, + {"configurable": {"thread_id": "interactive"}}, + ) + print(response["messages"][-1].content) + +# --------------------------------------------------------------------------- +# Main entry point +# --------------------------------------------------------------------------- +async def main(): + # Optional: load initial documents + # await load_documents_from_dir("./initial_docs") + await interactive_loop() if __name__ == "__main__": asyncio.run(main())