diff --git a/agent.py b/agent.py index 029650e..27a8c50 100644 --- a/agent.py +++ b/agent.py @@ -1,296 +1,51 @@ -""" -rag_agent.py — AI-агент с RAG-памятью на базе Ollama + Qdrant. +from langchain_ollama import ChatOllama +from langchain.agents import create_react_agent, AgentExecutor +from langchain_core.prompts import PromptTemplate +from tools import search_knowledge_base, add_to_knowledge_base -Стек: langchain-ollama, langchain-qdrant, langgraph, qdrant-client, rich -Установка: pip install langchain langchain-core langchain-ollama langchain-qdrant - langchain-text-splitters langgraph qdrant-client rich -""" +LLM_MODEL = "llama3" -import os -from pathlib import Path -from typing import List +SYSTEM_PROMPT = """You are a helpful AI assistant with access to a knowledge base. +Always use the knowledge base tools to search for relevant information before answering questions. +When you receive new information that should be remembered, add it to the knowledge base. -from langchain_ollama import OllamaEmbeddings, ChatOllama -from langchain_qdrant import QdrantVectorStore -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_core.tools import tool -from langchain_core.documents import Document -from langchain.agents import create_agent -from qdrant_client import QdrantClient as _QdrantClient -from qdrant_client.http.models import Distance, VectorParams +You have access to the following tools: -# --------------------------------------------------------------------------- -# Конфигурация -# --------------------------------------------------------------------------- +{tools} -QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost") -QDRANT_PORT = int(os.getenv("QDRANT_PORT", "6333")) -COLLECTION_NAME = os.getenv("QDRANT_COLLECTION", "rag_collection") -EMBEDDING_MODEL = os.getenv("EMBEDDING_MODEL", "nomic-embed-text") -LLM_MODEL = os.getenv("LLM_MODEL", "llama3") +Use the following format: -# --------------------------------------------------------------------------- -# Векторное хранилище (Qdrant + Ollama-эмбеддинги) -# --------------------------------------------------------------------------- +Question: the input question you must answer +Thought: you should always think about what to do +Action: the action to take, should be one of [{tool_names}] +Action Input: the input to the action +Observation: the result of the action +... (this Thought/Action/Action Input/Observation can repeat N times) +Thought: I now know the final answer +Final Answer: the final answer to the original input question -class RagVectorStore: - """Обёртка над QdrantVectorStore с эмбеддингами через Ollama.""" +Begin! - def __init__(self, collection_name: str = COLLECTION_NAME): - self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL) - - self._client = _QdrantClient(host=QDRANT_HOST, port=QDRANT_PORT) - - if not self._collection_exists(collection_name): - self._client.create_collection( - collection_name=collection_name, - vectors_config=VectorParams(size=768, distance=Distance.COSINE), - ) - - self.store = QdrantVectorStore( - client=self._client, - collection_name=collection_name, - embedding=self.embeddings, - ) - - def _collection_exists(self, name: str) -> bool: - try: - self._client.get_collection(name) - return True - except Exception: - return False - - def add_documents(self, docs: List[Document]) -> None: - self.store.add_documents(docs) - - def search(self, query: str, k: int = 5) -> List[tuple]: - # similarity_search_with_relevance_scores возвращает score в [0..1], - # где 1.0 = максимальная релевантность (исправлено с with_score) - results = self.store.similarity_search_with_relevance_scores(query, k=k) - return [(doc.page_content, score) for doc, score in results] +Question: {input} +Thought:{agent_scratchpad}""" -# --------------------------------------------------------------------------- -# Чанкинг -# --------------------------------------------------------------------------- - -def chunk_document(text: str, title: str) -> List[Document]: - """Разбивает текст на чанки с метаданными через RecursiveCharacterTextSplitter.""" - splitter = RecursiveCharacterTextSplitter( - chunk_size=1000, - chunk_overlap=200, - separators=["\n\n", "\n", " ", ""], +def create_rag_agent() -> AgentExecutor: + llm = ChatOllama(model=LLM_MODEL, temperature=0) + tools = [search_knowledge_base, add_to_knowledge_base] + prompt = PromptTemplate.from_template(SYSTEM_PROMPT) + agent = create_react_agent(llm=llm, tools=tools, prompt=prompt) + agent_executor = AgentExecutor( + agent=agent, + tools=tools, + verbose=True, + handle_parsing_errors=True, + max_iterations=10, ) - return [ - Document( - page_content=chunk, - metadata={"title": title, "source": title, "chunk_index": i}, - ) - for i, chunk in enumerate(splitter.split_text(text)) - ] + return agent_executor -# --------------------------------------------------------------------------- -# Синглтон хранилища (используется инструментами агента) -# --------------------------------------------------------------------------- - -rag_store = RagVectorStore() - - -# --------------------------------------------------------------------------- -# RAG-инструменты агента (@tool) -# --------------------------------------------------------------------------- - -@tool -def search_knowledge_base(query: str, max_results: int = 5) -> str: - """ - Выполняет семантический поиск в базе знаний. - - Args: - query: поисковый запрос на естественном языке. - max_results: максимальное количество результатов (по умолчанию 5). - - Returns: - Строка с найденными фрагментами и оценками релевантности. - """ - results = rag_store.search(query, k=max_results) - if not results: - return "В базе знаний ничего не найдено по данному запросу." - lines = [ - f"[{i}] (score={score:.3f}) {text}" - for i, (text, score) in enumerate(results, start=1) - ] - return "\n\n".join(lines) - - -@tool -def add_to_knowledge_base(content: str, title: str = "") -> str: - """ - Добавляет новый документ в базу знаний. - - Args: - content: текстовое содержимое документа. - title: название / источник документа (необязательно). - - Returns: - Сообщение об успехе с количеством созданных чанков. - """ - docs = chunk_document(content, title=title) - rag_store.add_documents(docs) - return ( - f"Документ «{title or 'без названия'}» успешно добавлен. " - f"Создано чанков: {len(docs)}." - ) - - -# --------------------------------------------------------------------------- -# Агент (create_agent из langchain.agents) -# --------------------------------------------------------------------------- - -llm = ChatOllama(model=LLM_MODEL, temperature=0.3) - -SYSTEM_PROMPT = """Ты — интеллектуальный AI-ассистент с доступом к локальной базе знаний. - -Правила работы: -1. ВСЕГДА начинай с поиска через search_knowledge_base перед ответом на вопрос. -2. Если найдена релевантная информация — используй её и ссылайся на источник. -3. Если пользователь просит сохранить информацию — используй add_to_knowledge_base. -4. Если база знаний пуста по теме — честно сообщи об этом и помоги по своим знаниям. -5. Отвечай на том же языке, на котором задан вопрос. -""" - -agent = create_agent( - model=llm, - tools=[search_knowledge_base, add_to_knowledge_base], - system_prompt=SYSTEM_PROMPT, -) - - -def run_agent(user_input: str) -> str: - """Запускает агента и возвращает финальный ответ.""" - result = agent.invoke({"messages": [("user", user_input)]}) - for msg in reversed(result["messages"]): - if msg.type == "ai" and msg.content: - return msg.content - return "Агент не вернул ответ." - - -# --------------------------------------------------------------------------- -# Загрузка документов из директории -# --------------------------------------------------------------------------- - -def load_documents_from_directory(directory: str) -> int: - """ - Загружает все .txt и .md файлы из папки в базу знаний. - Возвращает количество загруженных файлов. - """ - path = Path(directory) - count = 0 - for file_path in path.rglob("*"): - if file_path.suffix.lower() in {".txt", ".md"}: - text = file_path.read_text(encoding="utf-8", errors="ignore") - if text.strip(): - docs = chunk_document(text, title=file_path.name) - rag_store.add_documents(docs) - print(f" Загружен: {file_path.name} ({len(docs)} чанков)") - count += 1 - return count - - -# --------------------------------------------------------------------------- -# Интерактивный CLI-клиент -# --------------------------------------------------------------------------- - -if __name__ == "__main__": - from rich.console import Console - from rich.panel import Panel - from rich.markdown import Markdown - from rich.rule import Rule - - console = Console() - - console.print(Panel( - "[bold]RAG-агент[/bold] · Qdrant + Ollama\n\n" - "[cyan]/add <текст>[/cyan] — добавить текст в базу знаний\n" - "[cyan]/add-file <путь>[/cyan] — загрузить файл в базу знаний\n" - "[cyan]/search <запрос>[/cyan] — прямой поиск (без агента)\n" - "[cyan]/quit[/cyan] — выйти", - border_style="dim", - )) - - while True: - try: - user_input = input("\n> ").strip() - except (EOFError, KeyboardInterrupt): - console.print("\n[yellow]Выход.[/yellow]") - break - - if not user_input: - continue - - # /quit - if user_input.lower() in ("/quit", "/exit"): - console.print("[yellow]Выход.[/yellow]") - break - - # /add-file <путь> - elif user_input.lower().startswith("/add-file "): - file_path = Path(user_input[10:].strip()) - if file_path.is_file(): - text = file_path.read_text(encoding="utf-8", errors="ignore") - result = add_to_knowledge_base.invoke( - {"content": text, "title": file_path.name} - ) - console.print(f"[green]✓ {result}[/green]") - else: - console.print(f"[red]Файл не найден: {file_path}[/red]") - - # /add <текст> ← исправлено: теперь принимает текст напрямую, - # а не только путь к файлу - elif user_input.lower().startswith("/add "): - raw = user_input[5:].strip() - path = Path(raw) - if path.is_file(): - # если передан путь к существующему файлу — загружаем файл - text = path.read_text(encoding="utf-8", errors="ignore") - result = add_to_knowledge_base.invoke( - {"content": text, "title": path.name} - ) - else: - # иначе сохраняем как текст напрямую - result = add_to_knowledge_base.invoke( - {"content": raw, "title": "manual"} - ) - console.print(f"[green]✓ {result}[/green]") - - # /search <запрос> - elif user_input.lower().startswith("/search "): - query = user_input[8:].strip() - console.print(Rule("Результаты поиска")) - results = rag_store.search(query, k=5) - if not results: - console.print("[yellow]Ничего не найдено.[/yellow]") - else: - for i, (text, score) in enumerate(results, 1): - console.print(Panel( - text, - title=f"[cyan]#{i} score={score:.3f}[/cyan]", - border_style="dim", - )) - - # /search без аргумента - elif user_input.lower() == "/search": - console.print("[yellow]Укажите запрос: /search <текст>[/yellow]") - - # запрос к агенту - else: - console.print("[dim]Агент думает...[/dim]") - try: - answer = run_agent(user_input) - console.print(Panel( - Markdown(answer), - title="[green]Агент[/green]", - border_style="green", - )) - except Exception as e: - console.print(f"[red]Ошибка агента: {e}[/red]") \ No newline at end of file +def run_agent(query: str) -> str: + agent = create_rag_agent() + result = agent.invoke({"input": query}) + return result.get("output", "") \ No newline at end of file