246 lines
10 KiB
Python
246 lines
10 KiB
Python
from typing import TypedDict, List, Optional, Annotated
|
||
import operator
|
||
from langchain_community.llms import Ollama
|
||
from langchain_community.vectorstores import Qdrant
|
||
from langchain_community.document_loaders import DirectoryLoader, PyPDFLoader
|
||
from langgraph.graph import StateGraph, END
|
||
from langgraph.graph.message import add_messages
|
||
from langchain_core.messages import AIMessage, SystemMessage
|
||
from langchain.tools import tool
|
||
from local_knowledge_base import LocalKnowledgeBase
|
||
|
||
# from langchain.agents import AgentExecutor, create_react_agent
|
||
import numpy as np
|
||
|
||
|
||
# Определяем состояние агента
|
||
class AgentState(TypedDict):
|
||
messages: Annotated[List, add_messages] # История сообщений
|
||
knowledge_base: Optional[str] # Релевантные документы
|
||
current_step: str # Текущий шаг выполнения
|
||
needs_search: bool # Нужен ли поиск в RAG
|
||
final_answer: Optional[str] # Финальный ответ
|
||
|
||
|
||
# Инициализация LLM через Ollama
|
||
llm = Ollama(
|
||
model="qwen2.5-coder-16k:14b",
|
||
temperature=0.1, # Низкая температура для консистентности
|
||
num_predict=1024, # Максимальная длина ответа
|
||
)
|
||
|
||
|
||
# Создаем инструменты для агента
|
||
@tool
|
||
def search_knowledge_base(knowledge_base: LocalKnowledgeBase, query: str) -> str:
|
||
"""
|
||
🔍 Поиск информации в локальной базе знаний (база кода и документов).
|
||
|
||
🎯 КОГДА ИСПОЛЬЗОВАТЬ:
|
||
- Когда нужно найти определение класса, функции или структуры.
|
||
- Когда нужно найти реализацию алгоритма.
|
||
- Когда нужно найти описание ошибки или паттерна.
|
||
- Когда нужно найти документацию по проекту.
|
||
|
||
🚫 КОГДА НЕ ИСПОЛЬЗОВАТЬ:
|
||
- Для простых математических вычислений (используй calculate).
|
||
- Для общих вопросов, не связанных с загруженными файлами.
|
||
|
||
📝 ПРИМЕРЫ ЗАПРОСОВ:
|
||
- "Как работает класс Parser?"
|
||
- "Найди функцию обработки JSON"
|
||
- "Где реализован алгоритм сортировки?"
|
||
- "Объясни паттерн Singleton в этом проекте"
|
||
"""
|
||
|
||
# --- ЛОГИКА ПОИСКА ---
|
||
# Здесь должен быть реальный вызов к базе знаний.
|
||
# Пример реализации (нужно подключить к kb_instance):
|
||
#
|
||
results = knowledge_base.search(query, top_k=5)
|
||
if results:
|
||
context = "\n\n--- Найденные фрагменты ---\n\n"
|
||
for i, res in enumerate(results, 1):
|
||
context += f"[{i}] {res['text']}\n"
|
||
if len(res["text"]) > 500:
|
||
context += "..."
|
||
return context
|
||
else:
|
||
return "В базе знаний не найдено информации по запросу."
|
||
# return "Найденная информация из базы знаний"
|
||
|
||
|
||
@tool
|
||
def calculate(expression: str) -> str:
|
||
"""Выполнение математических вычислений"""
|
||
try:
|
||
result = eval(expression)
|
||
return f"Результат: {result}"
|
||
except:
|
||
return "Ошибка в выражении"
|
||
|
||
|
||
@tool
|
||
def web_search(query: str) -> str:
|
||
"""Поиск в интернете (если нужно)"""
|
||
# Можно подключить локальный поиск через DuckDuckGo
|
||
return "Результаты поиска из интернета"
|
||
|
||
|
||
# --- Глобальная инициализация базы знаний ---
|
||
# Мы создаем экземпляр базы знаний здесь, но он будет передан в граф
|
||
# В production лучше передавать как аргумент, но для простоты оставим здесь
|
||
# kb_instance = LocalKnowledgeBase(project_name="default")
|
||
|
||
|
||
# Создаем граф агента
|
||
def create_agent_graph(knowledge_base_instance: LocalKnowledgeBase):
|
||
kb = knowledge_base_instance
|
||
workflow = StateGraph(AgentState)
|
||
|
||
# Узел: анализ запроса
|
||
def analyze_query(state: AgentState):
|
||
messages = state["messages"]
|
||
last_message = messages[-1].content if messages else ""
|
||
|
||
print(f"\n🔍 АНАЛИЗ ЗАПРОСА: {last_message}")
|
||
|
||
# 1. Исключаем запросы, которые точно НЕ нужны RAG
|
||
non_rag_keywords = [
|
||
"посчитай",
|
||
"рассчитай",
|
||
"сколько",
|
||
"сколько же",
|
||
"простой пример",
|
||
"1+1",
|
||
"2*2",
|
||
]
|
||
if any(kw in last_message.lower() for kw in non_rag_keywords):
|
||
return {"needs_search": False, "current_step": "generating_answer"}
|
||
|
||
# 2. Ищем паттерны, связанные с кодом и проектом
|
||
# Используем регулярные выражения или простые проверки
|
||
code_patterns = [
|
||
r"class\s+\w+", # Поиск классов
|
||
r"function\s+\w+", # Поиск функций
|
||
r"namespace\s+\w+", # Поиск пространств имен
|
||
r"#include", # Поиск включений
|
||
r"enum\s+\w+", # Поиск enum
|
||
r"struct\s+\w+", # Поиск структур
|
||
]
|
||
|
||
# Если запрос содержит технические термины, скорее всего нужен RAG
|
||
has_code_keywords = any(pattern in last_message for pattern in code_patterns)
|
||
|
||
# 3. Явные запросы на поиск информации
|
||
search_keywords = [
|
||
"как работает",
|
||
"найди",
|
||
"где реализовано",
|
||
"ошибка",
|
||
"баг",
|
||
"документация",
|
||
"описание",
|
||
"реализация",
|
||
"класс",
|
||
"функция",
|
||
"алгоритм",
|
||
"паттерн",
|
||
"структура",
|
||
"метод",
|
||
]
|
||
|
||
needs_search = any(
|
||
keyword in last_message.lower() for keyword in search_keywords
|
||
)
|
||
|
||
# Комбинированное решение
|
||
needs_search = needs_search or has_code_keywords
|
||
|
||
print(f"🔍 Паттерны кода найдены: {has_code_keywords}")
|
||
print(f"🔍 Ключевые слова найдены: {needs_search}")
|
||
|
||
return {"needs_search": needs_search, "current_step": "analyzing_query"}
|
||
|
||
# Узел: поиск в RAG
|
||
def rag_search(state: AgentState):
|
||
if not state["needs_search"]:
|
||
return {"knowledge_base": None, "current_step": "generating_answer"}
|
||
|
||
# Получаем последний запрос
|
||
query = state["messages"][-1].content
|
||
|
||
# Здесь должен быть реальный поиск в Qdrant
|
||
# Пока заглушка
|
||
# results = ["Документ 1: Информация о...", "Документ 2: Данные по..."]
|
||
search_result = search_knowledge_base.invoke(
|
||
{"knowledge_base": kb, "query": query}
|
||
)
|
||
|
||
return {
|
||
"knowledge_base": search_result,
|
||
"current_step": "generating_answer",
|
||
}
|
||
|
||
# Узел: генерация ответа
|
||
def generate_answer(state: AgentState):
|
||
# Добавляем отладочный вывод
|
||
print("\n=== 🧠 ПРОМПТ ДЛЯ LLM (что видит модель) ===")
|
||
print(f"Доступные инструменты:")
|
||
print(f"1. search_knowledge_base: {search_knowledge_base.description}")
|
||
print(f"2. calculate: {calculate.description}")
|
||
print(f"3. web_search: {web_search.description}")
|
||
print(f"\nКонтекст из поиска: {state.get('knowledge_base', '')}")
|
||
print(f"=== Конец промпта ===\n")
|
||
|
||
messages = state["messages"]
|
||
knowledge = state.get("knowledge_base", "")
|
||
|
||
# Формируем промпт с контекстом
|
||
prompt = f"""Ты — интеллектуальный ассистент по C++ и архитектуре ПО.
|
||
|
||
Твоя задача: отвечать на вопросы пользователя, опираясь на загруженный код и документацию.
|
||
|
||
Контекст из базы знаний:
|
||
{knowledge}
|
||
|
||
История разговора:
|
||
{messages[-5:] if len(messages) > 5 else messages}
|
||
|
||
Текущий запрос: {messages[-1].content if messages else ''}
|
||
|
||
ИНСТРУКЦИИ:
|
||
1. Если контекст из базы знаний содержит информацию, релевантную запросу:
|
||
- Используй его для точного ответа.
|
||
- Цитируй фрагменты кода, если это нужно для объяснения.
|
||
- Объясняй код простым языком.
|
||
|
||
2. Если контекст пуст или не релевантен:
|
||
- Ответь что не знаешь
|
||
|
||
4. Если в найденных фрагментах есть несколько упоминаний одного и того же класса:
|
||
- Объясни каждый из них отдельно.
|
||
|
||
Ответь максимально полезно и точно:"""
|
||
|
||
response = llm.invoke(prompt)
|
||
|
||
return {
|
||
"final_answer": response,
|
||
"current_step": "completed",
|
||
"messages": messages + [AIMessage(content=response)],
|
||
}
|
||
|
||
# Добавляем узлы в граф
|
||
workflow.add_node("analyze", analyze_query)
|
||
workflow.add_node("search", rag_search)
|
||
workflow.add_node("generate", generate_answer)
|
||
|
||
# Определяем edges (переходы)
|
||
workflow.set_entry_point("analyze")
|
||
workflow.add_edge("analyze", "search")
|
||
workflow.add_edge("search", "generate")
|
||
workflow.add_edge("generate", END)
|
||
|
||
return workflow.compile()
|