From 01249a7f4c28c2985d9154d68458412be43d8563 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=94=D0=B0=D0=BD=D0=B8=D0=B8=D0=BB=20=D0=92=D0=B8=D0=BA?= =?UTF-8?q?=D1=82=D0=BE=D1=80=D0=BE=D0=B2?= Date: Thu, 2 Jul 2026 17:20:59 +0000 Subject: [PATCH] =?UTF-8?q?fix:=20main.py=20=E2=80=94=20=D0=90=D0=B3=D0=B5?= =?UTF-8?q?=D0=BD=D1=82=20=D1=81=20RAG-=D0=BF=D0=B0=D0=BC=D1=8F=D1=82?= =?UTF-8?q?=D1=8C=D1=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- main.py | 172 +++++++++++++++++++++++--------------------------------- 1 file changed, 69 insertions(+), 103 deletions(-) diff --git a/main.py b/main.py index 1b29b10..1904712 100644 --- a/main.py +++ b/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 qdrant_client import QdrantClient +from deepagents.backends import CompositeBackend, LocalShellBackend, FilesystemBackend -# 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 -qdrant_client = QdrantClient(url="http://localhost:6333") +# Инициализация QdrantVectorStore 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) - if not docs: - return "No results." - return "\n".join( - f"{i+1}. {doc.page_content[:200]}..." for i, doc in enumerate(docs) - ) + """Поиск в базе знаний.""" + docs = vector_store.similarity_search(query, k=max_results) + return "\n".join(d.page_content for d in docs) if docs else "No results." -# Backend for deepagents -backend = CompositeBackend( - [ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), - ] -) +@tool +def add_to_knowledge_base(content: str, title: str = "doc") -> str: + """Добавление документа в базу знаний.""" + vector_store.add_documents([Document(page_content=content, metadata={"title": title})]) + return f"Added: {title}" -# System prompt guiding the agent -system_prompt = ( - "You are a helpful agent with access to a knowledge base. " - "Use the provided tools to search and add information. " - "When searching, return concise results. " - "When adding, confirm the number of chunks added." -) +# Backend для deepagents +backend = CompositeBackend([ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), +]) -# Create the deep agent +# Создание агента agent = create_deep_agent( - model=chat, - tools=[add_to_knowledge_base, search_knowledge_base], + model=llm, + 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) -> None: - """Load all .txt files from dir_path into the knowledge base.""" +def load_documents_from_dir(dir_path: str): + """Загрузка всех .txt файлов из директории в векторную базу.""" for root, _, files in os.walk(dir_path): for file in files: if file.lower().endswith(".txt"): - path = os.path.join(root, file) - with open(path, "r", encoding="utf-8") as f: + file_path = os.path.join(root, file) + with open(file_path, "r", encoding="utf-8") as f: content = f.read() - title = os.path.splitext(file)[0] - add_to_knowledge_base(content, title) + chunks = splitter.split_text(content) + 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 - добавить документ") + print("/search - поиск в базе") + 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() - if user_input.lower() == "/quit": - print("Goodbye!") + user_input = input("> ").strip() + if not user_input: + continue + if user_input.startswith("/quit"): + print("Goodbye.") break - elif user_input.lower() == "/add": - title = input("Title: ").strip() - print("Enter content (end with a single line containing only 'END'):") - lines: List[str] = [] - while True: - line = input() - if line.strip() == "END": - break - lines.append(line) - content = "\n".join(lines) - message = f"Add the following content to knowledge base with title '{title}'." + elif user_input.startswith("/add"): + parts = user_input.split(maxsplit=1) + if len(parts) < 2: + print("Usage: /add ") + continue + file_path = parts[1] + if not os.path.isfile(file_path): + print(f"File not found: {file_path}") + continue + with open(file_path, "r", encoding="utf-8") as f: + 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": - query = input("Query: ").strip() - message = f"Search knowledge base for: {query}" + elif user_input.startswith("/search"): + parts = user_input.split(maxsplit=1) + if len(parts) < 2: + print("Usage: /search ") + 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 - 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 ") - 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() + print("Unknown command. Use /add, /search, /quit.") if __name__ == "__main__": asyncio.run(main()) \ No newline at end of file