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

main.py

@@ -1,151 +1,117 @@
-# DESIGN DECISION: Use OllamaEmbeddings and ChatOllama instead of OpenAI to satisfy assignment requirement of local LLM and embeddings via Ollama.
-# NECESSITY: Assignment explicitly requires local LLM and embeddings via Ollama; using OpenAI would violate constraints and introduce API keys.
-# OPTIMALITY: Ollama provides zero-cost inference, lower latency, and full data control; no external network calls.
-# ALTERNATIVES CONSIDERED: OpenRouter or OpenAI; rejected due to requirement of local models and cost.
-
 import os
-import sys
 import asyncio
-from typing import List
-
 from langchain_ollama import ChatOllama, OllamaEmbeddings
-from langchain_core.documents import Document
 from langchain_qdrant import QdrantVectorStore
 from langchain_text_splitters import RecursiveCharacterTextSplitter
+from langchain_core.documents import Document
 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 deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend
-from qdrant_client import QdrantClient
 
-# Initialize embeddings and chat models
+# Инициализация эмбеддингов и LLM через Ollama
 embeddings = OllamaEmbeddings(model="nomic-embed-text")
-chat = ChatOllama(model="llama3")
+llm = ChatOllama(model="llama3")
 
-# Initialize Qdrant client and vector store
+# Инициализация QdrantVectorStore
-qdrant_client = QdrantClient(url="http://localhost:6333")
 vector_store = QdrantVectorStore(
-    client=qdrant_client,
+    host="localhost",
+    port=6333,
     collection_name="knowledge",
     embedding_function=embeddings,
 )
 
-# Text splitter for chunking documents
+# Чанкинг текста
 splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
 
-# Tool: Add content to knowledge base
-@tool
-def add_to_knowledge_base(content: str, title: str = "doc") -> str:
-    """Add content to the knowledge base."""
-    chunks: List[str] = splitter.split_text(content)
-    docs: List[Document] = [
-        Document(page_content=chunk, metadata={"title": title}) for chunk in chunks
-    ]
-    vector_store.add_documents(docs)
-    return f"Added {len(docs)} chunks for {title}"
-
-# Tool: Search knowledge base
 @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(
-        f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs)
-    )
 
-# Backend for deepagents
+@tool
-backend = CompositeBackend(
+def add_to_knowledge_base(content: str, title: str = "doc") -> str:
-    [
+    """Добавление документа в базу знаний."""
-        LocalShellBackend(workspace_dir="./workspace"),
+    vector_store.add_documents([Document(page_content=content, metadata={"title": title})])
-        FilesystemBackend(),
+    return f"Added: {title}"
-    ]
-)
 
-# System prompt guiding the agent
+# Backend для deepagents
-system_prompt = (
+backend = CompositeBackend([
-    "You are a helpful agent with access to a knowledge base. "
+    LocalShellBackend(workspace_dir="./workspace"),
-    "Use the provided tools to search and add information. "
+    FilesystemBackend(),
-    "When searching, return concise results. "
+])
-    "When adding, confirm the number of chunks added."
-)
 
-# Create the deep agent
+# Создание агента
 agent = create_deep_agent(
-    model=chat,
+    model=llm,
-    tools=[add_to_knowledge_base, search_knowledge_base],
+    tools=[search_knowledge_base, add_to_knowledge_base],
     backend=backend,
-    system_prompt=system_prompt,
+    system_prompt="You are a helpful assistant. Use the provided tools to search and add knowledge.",
 )
 
-# Load documents from a directory into the knowledge base
+def load_documents_from_dir(dir_path: str):
-def load_documents_from_dir(dir_path: str) -> None:
+    """Загрузка всех .txt файлов из директории в векторную базу."""
-    """Load all .txt files from dir_path into the knowledge base."""
     for root, _, files in os.walk(dir_path):
         for file in files:
             if file.lower().endswith(".txt"):
-                path = os.path.join(root, file)
+                file_path = os.path.join(root, file)
-                with open(path, "r", encoding="utf-8") as f:
+                with open(file_path, "r", encoding="utf-8") as f:
                     content = f.read()
-                title = os.path.splitext(file)[0]
+                chunks = splitter.split_text(content)
-                add_to_knowledge_base(content, title)
+                docs = [Document(page_content=chunk, metadata={"title": file}) for chunk in chunks]
+                vector_store.add_documents(docs)
+
+async def main():
+    # Загрузка документов из папки data (если есть)
+    data_dir = "./data"
+    if os.path.isdir(data_dir):
+        load_documents_from_dir(data_dir)
+
+    print("RAG Agent ready. Commands:")
+    print("/add <file_path>   - добавить документ")
+    print("/search <query>    - поиск в базе")
+    print("/quit              - выйти")
 
-# Interactive CLI
-async def interactive_loop() -> None:
-    print("Welcome to the RAG agent CLI.")
-    print("Commands: /add, /search, /quit")
     while True:
-        user_input = input("\n> ").strip()
+        user_input = input("> ").strip()
-        if user_input.lower() == "/quit":
+        if not user_input:
-            print("Goodbye!")
+            continue
+        if user_input.startswith("/quit"):
+            print("Goodbye.")
             break
-        elif user_input.lower() == "/add":
+        elif user_input.startswith("/add"):
-            title = input("Title: ").strip()
+            parts = user_input.split(maxsplit=1)
-            print("Enter content (end with a single line containing only 'END'):")
+            if len(parts) < 2:
-            lines: List[str] = []
+                print("Usage: /add <file_path>")
-            while True:
+                continue
-                line = input()
+            file_path = parts[1]
-                if line.strip() == "END":
+            if not os.path.isfile(file_path):
-                    break
+                print(f"File not found: {file_path}")
-                lines.append(line)
+                continue
-            content = "\n".join(lines)
+            with open(file_path, "r", encoding="utf-8") as f:
-            message = f"Add the following content to knowledge base with title '{title}'."
+                content = f.read()
+            title = os.path.basename(file_path)
+            message = f"Add document titled {title} with content: {content}"
             result = await agent.ainvoke(
                 {"messages": [HumanMessage(content=message)]},
                 {"configurable": {"thread_id": "session-1"}},
             )
             print(result["messages"][-1].content)
-        elif user_input.lower() == "/search":
+        elif user_input.startswith("/search"):
-            query = input("Query: ").strip()
+            parts = user_input.split(maxsplit=1)
-            message = f"Search knowledge base for: {query}"
+            if len(parts) < 2:
+                print("Usage: /search <query>")
+                continue
+            query = parts[1]
+            message = f"Search the knowledge base for: {query}"
             result = await agent.ainvoke(
                 {"messages": [HumanMessage(content=message)]},
                 {"configurable": {"thread_id": "session-1"}},
             )
             print(result["messages"][-1].content)
         else:
-            # Treat as normal message
+            print("Unknown command. Use /add, /search, /quit.")
-            result = await agent.ainvoke(
-                {"messages": [HumanMessage(content=user_input)]},
-                {"configurable": {"thread_id": "session-1"}},
-            )
-            print(result["messages"][-1].content)
-
-async def main() -> None:
-    # Optional loading of documents via command line
-    if len(sys.argv) > 1 and sys.argv[1] == "--load-dir":
-        if len(sys.argv) < 3:
-            print("Usage: python main.py --load-dir <directory>")
-            return
-        dir_path = sys.argv[2]
-        if not os.path.isdir(dir_path):
-            print(f"Directory not found: {dir_path}")
-            return
-        print(f"Loading documents from {dir_path}...")
-        load_documents_from_dir(dir_path)
-        print("Loading complete.")
-    await interactive_loop()
 
 if __name__ == "__main__":
     asyncio.run(main())