From a00049957bcea412f4a8c0c4622a168a2b12e490 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 09:48: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 | 135 ++++++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 97 insertions(+), 38 deletions(-) diff --git a/main.py b/main.py index afa59fa..289b4d7 100644 --- a/main.py +++ b/main.py @@ -1,77 +1,136 @@ -import asyncio import os +import asyncio from pathlib import Path - -from langchain_core.messages import HumanMessage +from langchain_openai import ChatOpenAI, OpenAIEmbeddings +from langchain_core.documents import Document +from langchain_text_splitters import RecursiveCharacterTextSplitter +from langchain_qdrant import QdrantVectorStore from langchain.tools import tool from deepagents import create_deep_agent from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain_core.messages import HumanMessage -from utils import llm, vector_store, splitter -from langchain_core.documents import Document +# Configuration +OPENAI_API_KEY = os.getenv("OPENAI_API_KEY") +if not OPENAI_API_KEY: + raise RuntimeError("OPENAI_API_KEY not set in environment") +# LLM 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, +) + +# Embeddings via OpenRouter +embeddings = OpenAIEmbeddings( + model="text-embedding-3-small", + base_url="https://openrouter.ai/api/v1", + api_key=OPENAI_API_KEY, +) + +# Qdrant client and vector store +from qdrant_client import QdrantClient +qdrant_client = QdrantClient(url="http://localhost:6333") +collection_name = "knowledge_base" +vector_store = QdrantVectorStore( + client=qdrant_client, + collection_name=collection_name, + embedding_function=embeddings, +) + +# Text splitter +splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + +# Tool: search knowledge base @tool def search_knowledge_base(query: str, max_results: int = 3) -> str: """Search the knowledge base for relevant information.""" docs = vector_store.similarity_search(query, k=max_results) - return "\n".join(d.page_content for d in docs) if docs else "No results." + if not docs: + return "No results found." + return "\n\n".join(f"Title: {doc.metadata.get('title', 'unknown')}\n{doc.page_content}" for doc in docs) +# Tool: add to knowledge base @tool -def add_to_knowledge_base(content: str, title: str = "doc") -> str: +def add_to_knowledge_base(content: str, title: str = "document") -> 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 {len(docs)} chunks for {title}." + return f"Added {len(docs)} chunks from '{title}'." +# Backend for deepagents backend = CompositeBackend([ LocalShellBackend(workspace_dir="./workspace"), FilesystemBackend(), ]) +# Agent agent = create_deep_agent( model=llm, tools=[search_knowledge_base, add_to_knowledge_base], backend=backend, - system_prompt="You are a helpful assistant with access to a knowledge base. Use the provided tools to search and add information.", + system_prompt="You are a helpful knowledge assistant. Use the provided tools to search and add information.", ) -async def main(): - print("RAG Agent CLI. Commands: /add, /search, /quit") +# Helper: load documents from a directory +def load_documents_from_dir(directory: str): + dir_path = Path(directory) + if not dir_path.is_dir(): + raise ValueError(f"Directory {directory} does not exist.") + for file_path in dir_path.rglob("*"): + if file_path.is_file() and file_path.suffix.lower() in {".txt", ".md", ".py", ".json"}: + content = file_path.read_text(encoding="utf-8") + title = file_path.stem + add_to_knowledge_base(content, title) + +# Interactive CLI +async def interactive_loop(): + print("RAG Agent CLI. Commands: /add , /search , /quit") while True: user_input = input(">> ").strip() if not user_input: continue if user_input.lower() == "/quit": - print("Goodbye.") + print("Exiting.") break - if user_input.lower().startswith("/add"): - title = input("Title: ").strip() - print("Enter content (end with a single line containing only END):") - lines = [] - while True: - line = input() - if line.strip() == "END": - break - lines.append(line) - content = "\n".join(lines) - result = add_to_knowledge_base(content, title) + if user_input.startswith("/add "): + _, file_path = user_input.split(maxsplit=1) + try: + content = Path(file_path).read_text(encoding="utf-8") + title = Path(file_path).stem + result = add_to_knowledge_base(content, title) + print(result) + except Exception as e: + print(f"Error adding file: {e}") + continue + if user_input.startswith("/search "): + _, query = user_input.split(maxsplit=1) + result = search_knowledge_base(query, max_results=3) print(result) continue - if user_input.lower().startswith("/search"): - query = input("Query: ").strip() - max_results_str = input("Max results (default 3): ").strip() - max_results = int(max_results_str) if max_results_str.isdigit() else 3 - result = search_knowledge_base(query, max_results) - print("Search results:") - print(result) - continue - # Regular message to agent - response = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": "session-1"}}, - ) - print(response["messages"][-1].content) + # Treat as normal message to agent + try: + response = await agent.ainvoke( + {"messages": [HumanMessage(content=user_input)]}, + {"configurable": {"thread_id": "session-1"}}, + ) + print(response["messages"][-1].content) + except Exception as e: + print(f"Agent error: {e}") + +def main(): + # Optional: load initial docs from a folder + init_dir = os.getenv("INIT_DOCS_DIR") + if init_dir: + try: + load_documents_from_dir(init_dir) + print(f"Loaded documents from {init_dir}") + except Exception as e: + print(f"Failed to load initial docs: {e}") + asyncio.run(interactive_loop()) if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + main() \ No newline at end of file