diff --git a/main.py b/main.py index 77716b1..a4ab016 100644 --- a/main.py +++ b/main.py @@ -1,50 +1,47 @@ # main.py -# Полностью рабочий пример агента с RAG‑памятью на Qdrant и OpenRouter -# Использует deepagents, langchain‑openai, langchain‑qdrant, langchain‑core -# Запуск: python main.py +# Полностью рабочий пример агента с RAG‑памятью на базе Qdrant и Ollama. +# Используется LangChain 1.x, create_agent, инструменты @tool, и Ollama‑LLM/embeddings. +# +# Запуск: +# python main.py +# После запуска можно использовать команды: +# /add <content> – добавить документ +# /search <query> <max> – семантический поиск +# /quit – выйти +# +# Для загрузки документов из директории используйте функцию load_documents_from_dir. +#""" import os -import asyncio +import sys +import textwrap from pathlib import Path -from typing import List +from typing import List, Dict, Any -from langchain_openai import ChatOpenAI, OpenAIEmbeddings -from langchain_core.documents import Document +from langchain_ollama import ChatOllama, OllamaEmbeddings from langchain_qdrant import QdrantVectorStore from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain.tools import tool -from deepagents import create_deep_agent -from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend +from langchain.agents import create_agent, AgentExecutor, AgentToolkit, Tool from langchain_core.messages import HumanMessage # --------------------------------------------------------------------------- -# Конфигурация LLM и Embeddings +# Конфигурация # --------------------------------------------------------------------------- -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, -) - -embeddings = OpenAIEmbeddings( - model="text-embedding-3-small", - base_url="https://openrouter.ai/api/v1", - api_key=os.getenv("OPENAI_API_KEY"), -) +QDRANT_URL = os.getenv("QDRANT_URL", "http://localhost:6333") +QDRANT_COLLECTION = "knowledge_base" +EMBEDDING_MODEL = "nomic-embed-text" +LLM_MODEL = "llama3" # --------------------------------------------------------------------------- -# Qdrant клиент и коллекция +# Векторное хранилище # --------------------------------------------------------------------------- -# Предполагается, что Qdrant запущен локально на порту 6333 -qdrant_url = "http://localhost:6333" -collection_name = "knowledge_base" - +# Инициализируем эмбеддер и клиент Qdrant +embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) vector_store = QdrantVectorStore( - embeddings=embeddings, - url=qdrant_url, - collection_name=collection_name, - # Если коллекция не существует, она будет создана автоматически + url=QDRANT_URL, + collection_name=QDRANT_COLLECTION, + embedding=embeddings, ) # --------------------------------------------------------------------------- @@ -53,117 +50,131 @@ vector_store = QdrantVectorStore( text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) # --------------------------------------------------------------------------- -# Инструменты для агента +# Инструменты # --------------------------------------------------------------------------- -@tool +@tool("search_knowledge_base", "Semantic search in the knowledge base.") def search_knowledge_base(query: str, max_results: int = 5) -> str: - """Semantic search in the knowledge base. - Returns a formatted string with the top results. + """Return top‑k relevant documents for a query. + The function returns a formatted string with titles and snippets. """ results = vector_store.similarity_search_with_score(query, k=max_results) if not results: return "No relevant documents found." formatted = [] for doc, score in results: - formatted.append(f"Title: {doc.metadata.get('title', 'Untitled')}\nScore: {score:.4f}\nContent: {doc.page_content[:200]}...\n") + title = doc.metadata.get("title", "Untitled") + snippet = doc.page_content[:200].replace("\n", " ") + formatted.append(f"{title} (score: {score:.3f}): {snippet}...") return "\n".join(formatted) -@tool +@tool("add_to_knowledge_base", "Add a document to the knowledge base.") def add_to_knowledge_base(content: str, title: str) -> str: - """Add a new document to the knowledge base. - The content is split into chunks, embedded and stored. + """Chunk the content, embed, and store in Qdrant. + Returns a confirmation message. """ - # Split into chunks chunks = text_splitter.split_text(content) - docs: List[Document] = [] + docs = [] for i, chunk in enumerate(chunks): - docs.append(Document(page_content=chunk, metadata={"title": title, "chunk_index": i})) - # Add to vector store - vector_store.add_documents(docs) - return f"Added {len(docs)} chunks for document '{title}'." - -# --------------------------------------------------------------------------- -# Backend для deepagents -# --------------------------------------------------------------------------- -backend = CompositeBackend( - default=LocalShellBackend(root_dir="./workspace", virtual_mode=True, inherit_env=True), - routes={}, -) - -# --------------------------------------------------------------------------- -# Создание агента -# --------------------------------------------------------------------------- -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." -) - -# --------------------------------------------------------------------------- -# Инициализация: загрузка документов из директории -# --------------------------------------------------------------------------- -DOCS_DIR = Path("./docs") - -async def load_documents_from_dir(directory: Path): - if not directory.exists(): - return - for file_path in directory.rglob("*.txt"): - content = file_path.read_text(encoding="utf-8") - title = file_path.stem - await agent.ainvoke( - {"messages": [HumanMessage(content=f"/add {title}")], "content": content}, - {"configurable": {"thread_id": "init-session"}}, + docs.append( + { + "page_content": chunk, + "metadata": {"title": title, "chunk_index": i}, + } ) + vector_store.add_documents(docs) + return f"Added {len(chunks)} chunks of '{title}' to the knowledge base." # --------------------------------------------------------------------------- -# Интерактивный клиент +# Агент # --------------------------------------------------------------------------- -async def interactive_loop(): - print("Welcome to the RAG agent. Commands: /add <title>, /search <query>, /quit") - thread_id = "interactive-session" +# Создаём LLM +llm = ChatOllama(model=LLM_MODEL, temperature=0.2) + +# Список инструментов +tools = [search_knowledge_base, add_to_knowledge_base] + +# Создаём агент +agent = create_agent( + llm=llm, + tools=tools, + system_message="You are an assistant that can search and add documents to a local knowledge base. Use the provided tools.", + verbose=True, +) + +# Обёртка для выполнения +agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True) + +# --------------------------------------------------------------------------- +# Загрузка документов из директории +# --------------------------------------------------------------------------- + +def load_documents_from_dir(directory: str) -> None: + """Load all .txt files from a directory into the knowledge base. + Each file becomes a separate document with its filename as title. + """ + path = Path(directory) + if not path.is_dir(): + print(f"Directory {directory} does not exist.") + return + for file in path.glob("*.txt"): + title = file.stem + content = file.read_text(encoding="utf-8") + print(f"Adding {title}...", end=" ") + result = add_to_knowledge_base(content, title) + print(result) + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + +def main(): + # Если пользователь передал путь к директории, загрузим документы + if len(sys.argv) > 1: + load_documents_from_dir(sys.argv[1]) + + print("\n--- RAG Agent CLI ---") + print("Commands:") + print(" /add <title> <content> – add a document") + print(" /search <query> <max> – search knowledge base") + print(" /quit – exit") + while True: - user_input = input("> ") - if user_input.strip() == "/quit": + try: + user_input = input("\n> ") + except (EOFError, KeyboardInterrupt): + print("\nExiting.") + break + + if not user_input.strip(): + continue + + if user_input.startswith("/quit"): print("Goodbye!") break - if user_input.startswith("/add "): - title = user_input[5:].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) - await agent.ainvoke( - {"messages": [HumanMessage(content=f"/add {title}")], "content": content}, - {"configurable": {"thread_id": thread_id}}, - ) - print(f"Document '{title}' added.") - elif user_input.startswith("/search "): - query = user_input[8:].strip() - result = await agent.ainvoke( - {"messages": [HumanMessage(content=f"/search {query}")]}, - {"configurable": {"thread_id": thread_id}}, - ) - print(result["messages"][-1].content) - else: - # обычный запрос к LLM - result = await agent.ainvoke( - {"messages": [HumanMessage(content=user_input)]}, - {"configurable": {"thread_id": thread_id}}, - ) - print(result["messages"][-1].content) -# --------------------------------------------------------------------------- -# Основной запуск -# --------------------------------------------------------------------------- -async def main(): - # Загрузим документы из каталога docs при старте - await load_documents_from_dir(DOCS_DIR) - await interactive_loop() + if user_input.startswith("/add"): + parts = user_input.split(maxsplit=2) + if len(parts) < 3: + print("Usage: /add <title> <content>") + continue + title, content = parts[1], parts[2] + print(add_to_knowledge_base(content, title)) + continue + + if user_input.startswith("/search"): + parts = user_input.split(maxsplit=2) + if len(parts) < 2: + print("Usage: /search <query> [max_results]") + continue + query = parts[1] + max_results = int(parts[2]) if len(parts) > 2 else 5 + print(search_knowledge_base(query, max_results)) + continue + + # Любой другой ввод – передаём агенту + response = agent_executor.invoke({"input": user_input}) + print(response.get("output", "")) + if __name__ == "__main__": - asyncio.run(main()) + main()