From d04acc7a69cc0b8a2a827f863af6bbf189770b2e 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:08:27 +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 | 153 ++++++++++++++++++++++++++++++++++++++++++++++++-------- 1 file changed, 131 insertions(+), 22 deletions(-) diff --git a/main.py b/main.py index 087f7ef..1b29b10 100644 --- a/main.py +++ b/main.py @@ -1,42 +1,151 @@ -import asyncio -from langchain_core.messages import HumanMessage -from rag_agent import agent, add_to_knowledge_base, search_knowledge_base +# 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. -async def interactive_loop(): - print("Welcome to the RAG agent. Type /add to add a document, /search to query, /quit to exit.") - thread_id = "session-1" +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.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 + +# Initialize embeddings and chat models +embeddings = OllamaEmbeddings(model="nomic-embed-text") +chat = ChatOllama(model="llama3") + +# Initialize Qdrant client and vector store +qdrant_client = QdrantClient(url="http://localhost:6333") +vector_store = QdrantVectorStore( + client=qdrant_client, + 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) + ) + +# Backend for deepagents +backend = CompositeBackend( + [ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), + ] +) + +# 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." +) + +# Create the deep agent +agent = create_deep_agent( + model=chat, + tools=[add_to_knowledge_base, search_knowledge_base], + backend=backend, + system_prompt=system_prompt, +) + +# 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.""" + 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: + content = f.read() + title = os.path.splitext(file)[0] + add_to_knowledge_base(content, title) + +# Interactive CLI +async def interactive_loop() -> None: + print("Welcome to the RAG agent CLI.") + print("Commands: /add, /search, /quit") while True: - user_input = input(">> ").strip() + user_input = input("\n> ").strip() if user_input.lower() == "/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 = [] + 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) + message = f"Add the following content to knowledge base with title '{title}'." + 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() - result = search_knowledge_base(query) - print("Search results:") - print(result) - else: - messages = [HumanMessage(content=user_input)] - response = await agent.ainvoke( - {"messages": messages}, - {"configurable": {"thread_id": thread_id}}, + message = f"Search knowledge base for: {query}" + result = await agent.ainvoke( + {"messages": [HumanMessage(content=message)]}, + {"configurable": {"thread_id": "session-1"}}, ) - print(response["messages"][-1].content) + 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) -def main(): - asyncio.run(interactive_loop()) +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() if __name__ == "__main__": - main() \ No newline at end of file + asyncio.run(main()) \ No newline at end of file