mariahberdnikova/task-6a02e23d-agent-s-rag
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Обновить agent.py

Browse Source
This commit is contained in:
Мария Бердникова
2026-05-28 13:04:50 +00:00
parent d44911dded
commit 7c76c16bc5
1 changed files with 38 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent.py
+38 -283
View File
@@ -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)
Question: {input}
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),
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,
)
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]
return agent_executor
# ---------------------------------------------------------------------------
# Чанкинг
# ---------------------------------------------------------------------------
def chunk_document(text: str, title: str) -> List[Document]:
"""Разбивает текст на чанки с метаданными через RecursiveCharacterTextSplitter."""
splitter = RecursiveCharacterTextSplitter(
chunk_size=1000,
chunk_overlap=200,
separators=["\n\n", "\n", " ", ""],
)
return [
Document(
page_content=chunk,
metadata={"title": title, "source": title, "chunk_index": i},
)
for i, chunk in enumerate(splitter.split_text(text))
]
# ---------------------------------------------------------------------------
# Синглтон хранилища (используется инструментами агента)
# ---------------------------------------------------------------------------
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]")
def run_agent(query: str) -> str:
agent = create_rag_agent()
result = agent.invoke({"input": query})
return result.get("output", "")