From a71260af3c9363fc0a393e76c8437da5862ae9f9 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 08:59:12 +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 | 162 ++++++++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 128 insertions(+), 34 deletions(-) diff --git a/main.py b/main.py index 75fb450..96f4768 100644 --- a/main.py +++ b/main.py @@ -1,58 +1,152 @@ import os +import sys import asyncio -from dotenv import load_dotenv -from langchain_openai import ChatOpenAI +from pathlib import Path +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 tools import search_knowledge_base, add_to_knowledge_base +from qdrant_client import QdrantClient -load_dotenv() +# Инициализация эмбеддингов и LLM через Ollama +embeddings = OllamaEmbeddings(model="nomic-embed-text") +llm = Ollama(model="llama3") -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, +# Инициализация Qdrant +client = QdrantClient(host="localhost", port=6333) +collection_name = "knowledge" +vector_store = QdrantVectorStore( + client=client, + collection_name=collection_name, + embedding=embeddings, ) -backend = CompositeBackend([ - LocalShellBackend(workspace_dir="./workspace"), - FilesystemBackend(), -]) +# Чанкинг +splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200) + +# Инструмент: поиск в базе знаний +@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(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)." + +# Backend для deepagents +backend = CompositeBackend( + [ + LocalShellBackend(workspace_dir="./workspace"), + FilesystemBackend(), + ] +) + +system_prompt = ( + "You are a helpful agent with access to a knowledge base. " + "Use the tools to search and add information. " + "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="You are a helpful knowledge assistant. Use the tools to search and add documents.", + system_prompt=system_prompt, ) -async def interactive_loop(): - thread_id = "interactive-session" - print("Welcome to RAG Agent. Commands: /add <content>, /search <query>, /quit") +# Загрузка документов из директории +def load_documents(dir_path: str) -> None: + """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: - user_input = input(">> ") - if user_input.strip() == "/quit": + try: + user_input = input("\n> ").strip() + except (EOFError, KeyboardInterrupt): + print("\nExiting.") + break + + if not user_input: + continue + + if user_input.lower() == "/quit": print("Goodbye.") break - if user_input.startswith("/add"): - try: - _, title, content = user_input.split(" ", 2) - except ValueError: - print("Usage: /add <title> <content>") + + if 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) + 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 - message = HumanMessage(content=f"Add document titled '{title}' with content: {content}") - elif user_input.startswith("/search"): - query = user_input[len("/search"):].strip() - message = HumanMessage(content=f"Search knowledge base for: {query}") - else: - message = HumanMessage(content=user_input) - result = await agent.ainvoke( - {"messages": [message]}, + result = search_knowledge_base(query) + print("\nSearch results:") + print(result) + continue + + # Any other message is sent to the agent + response = await agent.ainvoke( + {"messages": [HumanMessage(content=user_input)]}, {"configurable": {"thread_id": thread_id}}, ) - print(result["messages"][-1].content) + agent_reply = response["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__": - asyncio.run(interactive_loop()) \ No newline at end of file + main() \ No newline at end of file