Обновить agent.py
This commit is contained in:
@@ -1,296 +1,51 @@
|
|||||||
"""
|
from langchain_ollama import ChatOllama
|
||||||
rag_agent.py — AI-агент с RAG-памятью на базе Ollama + Qdrant.
|
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
|
LLM_MODEL = "llama3"
|
||||||
Установка: pip install langchain langchain-core langchain-ollama langchain-qdrant
|
|
||||||
langchain-text-splitters langgraph qdrant-client rich
|
|
||||||
"""
|
|
||||||
|
|
||||||
import os
|
SYSTEM_PROMPT = """You are a helpful AI assistant with access to a knowledge base.
|
||||||
from pathlib import Path
|
Always use the knowledge base tools to search for relevant information before answering questions.
|
||||||
from typing import List
|
When you receive new information that should be remembered, add it to the knowledge base.
|
||||||
|
|
||||||
from langchain_ollama import OllamaEmbeddings, ChatOllama
|
You have access to the following tools:
|
||||||
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
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
{tools}
|
||||||
# Конфигурация
|
|
||||||
# ---------------------------------------------------------------------------
|
|
||||||
|
|
||||||
QDRANT_HOST = os.getenv("QDRANT_HOST", "localhost")
|
Use the following format:
|
||||||
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")
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
Question: the input question you must answer
|
||||||
# Векторное хранилище (Qdrant + Ollama-эмбеддинги)
|
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:
|
Begin!
|
||||||
"""Обёртка над QdrantVectorStore с эмбеддингами через Ollama."""
|
|
||||||
|
|
||||||
def __init__(self, collection_name: str = COLLECTION_NAME):
|
Question: {input}
|
||||||
self.embeddings = OllamaEmbeddings(model=EMBEDDING_MODEL)
|
Thought:{agent_scratchpad}"""
|
||||||
|
|
||||||
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]
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
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)
|
||||||
def chunk_document(text: str, title: str) -> List[Document]:
|
agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
|
||||||
"""Разбивает текст на чанки с метаданными через RecursiveCharacterTextSplitter."""
|
agent_executor = AgentExecutor(
|
||||||
splitter = RecursiveCharacterTextSplitter(
|
agent=agent,
|
||||||
chunk_size=1000,
|
tools=tools,
|
||||||
chunk_overlap=200,
|
verbose=True,
|
||||||
separators=["\n\n", "\n", " ", ""],
|
handle_parsing_errors=True,
|
||||||
|
max_iterations=10,
|
||||||
)
|
)
|
||||||
return [
|
return agent_executor
|
||||||
Document(
|
|
||||||
page_content=chunk,
|
|
||||||
metadata={"title": title, "source": title, "chunk_index": i},
|
|
||||||
)
|
|
||||||
for i, chunk in enumerate(splitter.split_text(text))
|
|
||||||
]
|
|
||||||
|
|
||||||
|
|
||||||
# ---------------------------------------------------------------------------
|
def run_agent(query: str) -> str:
|
||||||
# Синглтон хранилища (используется инструментами агента)
|
agent = create_rag_agent()
|
||||||
# ---------------------------------------------------------------------------
|
result = agent.invoke({"input": query})
|
||||||
|
return result.get("output", "")
|
||||||
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]")
|
|
||||||
Reference in New Issue
Block a user