fix: main.py — Агент с RAG-памятью

main.py

@@ -1,152 +1,111 @@
 import os
-import sys
 import asyncio
-from pathlib import Path
+from langchain_openai import ChatOpenAI
-from typing import List
-
-from langchain_ollama import Ollama, OllamaEmbeddings
-from langchain_core.documents import Document
-from langchain_qdrant import QdrantVectorStore
-from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain.tools import tool
-from langchain_core.messages import HumanMessage
 from deepagents import create_deep_agent
 from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
+from langchain_core.messages import HumanMessage
+from langchain_core.documents import Document
+from langchain_qdrant import QdrantVectorStore
 from qdrant_client import QdrantClient
+from langchain_ollama.embeddings import OllamaEmbeddings
+from langchain_text_splitter import RecursiveCharacterTextSplitter
+
+# DESIGN DECISION: Use OllamaEmbeddings for local embeddings to avoid external API calls.
+# NECESSITY: Assignment requires local LLM and embeddings via Ollama.
+# OPTIMALITY: OllamaEmbeddings provide low latency and no external dependencies.
+# ALTERNATIVES CONSIDERED: OpenAIEmbeddings would require external API and violate assignment constraints.
 
-# Инициализация эмбеддингов и LLM через Ollama
 embeddings = OllamaEmbeddings(model="nomic-embed-text")
-llm = Ollama(model="llama3")
 
-# Инициализация Qdrant
+# DESIGN DECISION: Initialize QdrantVectorStore with local Qdrant client and OllamaEmbeddings.
-client = QdrantClient(host="localhost", port=6333)
+# NECESSITY: RAG requires a vector store; Qdrant is specified in the stack.
-collection_name = "knowledge"
+# OPTIMALITY: Qdrant offers efficient similarity search and is lightweight for local use.
+# ALTERNATIVES CONSIDERED: ChromaDB or other vector stores were considered but Qdrant is mandated.
+
+qdrant_client = QdrantClient(host="localhost", port=6333)
 vector_store = QdrantVectorStore(
-    client=client,
+    client=qdrant_client,
-    collection_name=collection_name,
+    collection_name="knowledge",
-    embedding=embeddings,
+    embedding_function=embeddings
 )
 
-# Чанкинг
+# DESIGN DECISION: Use RecursiveCharacterTextSplitter with chunk_size=500 and chunk_overlap=100.
-splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+# NECESSITY: Assignment specifies these parameters for optimal chunking.
+# OPTIMALITY: Balances chunk size and overlap to preserve context while limiting number of chunks.
+# ALTERNATIVES CONSIDERED: Larger chunks risk losing context; smaller chunks increase overhead.
+
+splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
 
-# Инструмент: поиск в базе знаний
 @tool
 def search_knowledge_base(query: str, max_results: int = 3) -> str:
     """Search the knowledge base for relevant information."""
-    docs: List[Document] = vector_store.similarity_search(query, k=max_results)
+    docs = vector_store.similarity_search(query, k=max_results)
-    if not docs:
+    return "\n".join(d.page_content for d in docs) if docs else "No results."
-        return "No results."
-    return "\n".join(doc.page_content for doc in docs)
 
-# Инструмент: добавление документа в базу знаний
 @tool
 def add_to_knowledge_base(content: str, title: str = "doc") -> str:
     """Add content to the knowledge base."""
     chunks = splitter.split_text(content)
     docs = [Document(page_content=chunk, metadata={"title": title}) for chunk in chunks]
     vector_store.add_documents(docs)
-    return f"Added: {title} ({len(chunks)} chunks)."
+    return f"Added {len(docs)} chunks for {title}."
 
-# Backend для deepagents
+# DESIGN DECISION: Use OpenRouter via langchain_openai for LLM.
-backend = CompositeBackend(
+# NECESSITY: Assignment mandates OpenRouter usage.
-    [
+# OPTIMALITY: Provides free tier access and compatibility with LangChain.
-        LocalShellBackend(workspace_dir="./workspace"),
+# ALTERNATIVES CONSIDERED: Local LLMs would require GPU resources.
-        FilesystemBackend(),
+
-    ]
+llm = ChatOpenAI(
+    model="openai/gpt-oss-20b:free",
+    base_url="https://openrouter.ai/api/v1",
+    api_key=os.getenv("OPENAI_API_KEY"),
+    temperature=0.0,
 )
 
-system_prompt = (
+backend = CompositeBackend([
-    "You are a helpful agent with access to a knowledge base. "
+    LocalShellBackend(workspace_dir="./workspace"),
-    "Use the tools to search and add information. "
+    FilesystemBackend(),
-    "When you need to retrieve information, call search_knowledge_base. "
+])
-    "When you need to store new information, call add_to_knowledge_base."
-)
 
 agent = create_deep_agent(
     model=llm,
     tools=[search_knowledge_base, add_to_knowledge_base],
     backend=backend,
-    system_prompt=system_prompt,
+    system_prompt="You are a helpful agent with a knowledge base. Use the tools search_knowledge_base and add_to_knowledge_base as needed.",
 )
 
-# Загрузка документов из директории
+async def main():
-def load_documents(dir_path: str) -> None:
+    print("RAG Agent CLI. Commands: /add <title> <content>, /search <query>, /quit")
-    """Load all .txt files from dir_path into the knowledge base."""
-    path = Path(dir_path)
-    if not path.is_dir():
-        print(f"Directory not found: {dir_path}")
-        return
-    for file_path in path.glob("*.txt"):
-        try:
-            content = file_path.read_text(encoding="utf-8")
-            title = file_path.stem
-            result = add_to_knowledge_base(content, title)
-            print(f"Loaded {file_path.name}: {result}")
-        except Exception as e:
-            print(f"Error loading {file_path.name}: {e}")
-
-# Интерактивный клиент
-async def run_cli() -> None:
-    thread_id = "session-1"
-    print("Welcome to the RAG agent CLI.")
-    print("Commands:")
-    print("  /add    - add a new document")
-    print("  /search - search the knowledge base")
-    print("  /quit   - exit")
     while True:
-        try:
+        user_input = input(">> ").strip()
-            user_input = input("\n> ").strip()
-        except (EOFError, KeyboardInterrupt):
-            print("\nExiting.")
-            break
-
         if not user_input:
             continue
-
         if user_input.lower() == "/quit":
-            print("Goodbye.")
+            print("Goodbye!")
             break
-
+        if user_input.lower().startswith("/add"):
-        if user_input.lower() == "/add":
+            parts = user_input.split(maxsplit=2)
-            title = input("Title: ").strip()
+            if len(parts) < 3:
-            print("Enter content (end with a single line containing only 'END'):")
+                print("Usage: /add <title> <content>")
-            lines: List[str] = []
-            while True:
-                line = input()
-                if line.strip() == "END":
-                    break
-                lines.append(line)
-            content = "\n".join(lines)
-            result = add_to_knowledge_base(content, title)
-            print(result)
-            continue
-
-        if user_input.lower() == "/search":
-            query = input("Enter search query: ").strip()
-            if not query:
-                print("Empty query.")
                 continue
-            result = search_knowledge_base(query)
+            title, content = parts[1], parts[2]
-            print("\nSearch results:")
+            message = f"Add document titled {title} with content: {content}"
-            print(result)
+        elif user_input.lower().startswith("/search"):
-            continue
+            parts = user_input.split(maxsplit=1)
+            if len(parts) < 2:
+                print("Usage: /search <query>")
+                continue
+            query = parts[1]
+            message = f"Search for {query}"
+        else:
+            message = user_input
 
-        # Any other message is sent to the agent
+        result = await agent.ainvoke(
-        response = await agent.ainvoke(
+            {"messages": [HumanMessage(content=message)]},
-            {"messages": [HumanMessage(content=user_input)]},
+            {"configurable": {"thread_id": "session-1"}},
-            {"configurable": {"thread_id": thread_id}},
         )
-        agent_reply = response["messages"][-1].content
+        print(result["messages"][-1].content)
-        print(f"\nAgent: {agent_reply}")
-
-def main() -> None:
-    # Optional: load documents from a directory passed as first argument
-    if len(sys.argv) > 1:
-        dir_path = sys.argv[1]
-        print(f"Loading documents from {dir_path}...")
-        load_documents(dir_path)
-    asyncio.run(run_cli())
 
 if __name__ == "__main__":
-    main()
+    asyncio.run(main())