analyzer/tests/tools.py - comb - obuxxxov/CodeForge

analyzer/tests/tools.py
264 строки10 KB
ObliInSe
start
18 авг 2025, 00:31
18 авг 2025, 00:310671f4e
from __future__ import annotations

import operator
import os
import re
from typing import Any, Dict, List, Optional

# Предпочитаем стандартную библиотеку typing; fallback на typing_extensions если нужно
try:
	from typing import Annotated, TypedDict
except ImportError:  # pragma: no cover
	from typing_extensions import Annotated, TypedDict

# Пытаемся импортировать классы сообщений LangChain; fallback на простые заглушки если недоступны
try:
	from langchain_core.messages import AIMessage, AnyMessage, HumanMessage
except ImportError:  # pragma: no cover
	class _BaseMessage:  # минимальная заглушка с .content
		def __init__(self, content: str):
			self.content = content
	AIMessage = _BaseMessage
	HumanMessage = _BaseMessage
	AnyMessage = _BaseMessage

from langgraph.graph import END, StateGraph

# MCP интеграция
try:
	import asyncio
	from mcp import ClientSession, StdioServerParameters
	from mcp.types import Tool, TextContent, ImageContent
	
	MCP_AVAILABLE = True
	print("✅ MCP протокол доступен")
except ImportError:
	MCP_AVAILABLE = False
	print("⚠️ MCP не установлен. Установите: pip install mcp")

# Интеграция с GigaChat
try:
	from dotenv import load_dotenv
	from gigachat import GigaChat
	
	# Загружаем переменные окружения
	load_dotenv()
	
	# Получаем credentials из переменных окружения
	credentials = os.getenv('GIGACHAT_CREDENTIALS')
	
	if not credentials:
		raise ValueError("GIGACHAT_CREDENTIALS не найден в переменных окружения. Пожалуйста, создайте .env файл с вашими credentials.")
	
	giga = GigaChat(
		credentials=credentials,
		verify_ssl_certs=False,
	)
	print("✅ GigaChat подключен")
		
except ImportError:
	raise ImportError("GigaChat не установлен. Установите: pip install gigachat")
except Exception as e:
	raise Exception(f"Ошибка подключения к GigaChat: {e}")


class AgentState(TypedDict):
	"""Состояние для минимального агента.

	- messages: история сообщений, агрегируется через добавление
	- tool_result: опциональная строка с результатом инструмента для текущего хода
	- mcp_tools: список доступных MCP инструментов
	"""

	messages: Annotated[List[AnyMessage], operator.add]
	tool_result: Optional[str]
	mcp_tools: Optional[List[Tool]]


def _classify_intent(text: str) -> str:
	"""Очень простой эвристический классификатор: решает использовать ли инструмент.

	Возвращает одно из: "tool" | "mcp" | "respond".
	"""
	# Если есть хотя бы два числа или явные математические глаголы, направляем к математическому инструменту
	numbers = re.findall(r"[-+]?\d*\.?\d+", text)
	if len(numbers) >= 2:
		return "tool"
	if re.search(r"\b(sum|add|plus|склади|сложи|сложить|сумма|прибавь|посчитай|умножь|multiply|product)\b", text, re.IGNORECASE):
		return "tool"
	
	# Если есть слова связанные с поиском/информацией, используем MCP
	if re.search(r"\b(найди|поиск|search|find|погода|weather|курс|rate|новости|news)\b", text, re.IGNORECASE):
		return "mcp"
	
	return "respond"


def _passthrough_node(state: AgentState) -> Dict[str, Any]:
	# Узел-заглушка, используемый как точка входа
	return {}


def _decide_next(state: AgentState) -> str:
	last_user = None
	for msg in reversed(state.get("messages", [])):
		if isinstance(msg, HumanMessage):
			last_user = msg.content
			break
	last_user = last_user or ""
	decision = _classify_intent(str(last_user))
	print(f"🔍 Решение: '{last_user}' -> {decision}")
	return decision


def _math_tool_node(state: AgentState) -> Dict[str, Any]:
	"""Крошечный встроенный "калькулятор" который суммирует или умножает числа из текста.

	Эвристики:
	- Если текст содержит слова для умножения, вычисляем произведение; иначе сумму.
	- Извлекает все числа (целые/дробные) из последнего сообщения пользователя.
	"""
	print("🔧 Выполняется МАТЕМАТИЧЕСКИЙ ИНСТРУМЕНТ")
	last_user = None
	for msg in reversed(state.get("messages", [])):
		if isinstance(msg, HumanMessage):
			last_user = msg.content
			break

	text = str(last_user or "")
	values = [float(x) for x in re.findall(r"[-+]?\d*\.?\d+", text)]
	if not values:
		return {"tool_result": "Не удалось найти числа для вычислений."}

	is_multiply = bool(re.search(r"\b(умнож|multiply|product|произвед)\w*\b", text, re.IGNORECASE))
	if is_multiply:
		result = 1.0
		for v in values:
			result *= v
		op = "произведение"
	else:
		result = sum(values)
		op = "сумма"

	# Сохраняем результат инструмента; узел respond сформирует финальное сообщение
	tool_result = f"{op} чисел {values} = {result}"
	print(f"🔧 Результат инструмента: {tool_result}")
	return {"tool_result": tool_result}


async def _mcp_tool_node(state: AgentState) -> Dict[str, Any]:
	"""Узел для вызова MCP инструментов.

	Поддерживает внешние инструменты через MCP протокол.
	"""
	print("🌐 Выполняется MCP ИНСТРУМЕНТ")
	
	if not MCP_AVAILABLE:
		return {"tool_result": "MCP не доступен. Установите: pip install mcp"}
	
	last_user = None
	for msg in reversed(state.get("messages", [])):
		if isinstance(msg, HumanMessage):
			last_user = msg.content
			break

	text = str(last_user or "")
	
	# Простая заглушка MCP инструмента для демонстрации
	# В реальности здесь был бы вызов внешнего MCP сервера
	if "погода" in text.lower() or "weather" in text.lower():
		tool_result = "MCP инструмент 'weather': Температура 22°C, солнечно, влажность 65%"
	elif "курс" in text.lower() or "rate" in text.lower():
		tool_result = "MCP инструмент 'currency': USD/RUB = 95.42, EUR/RUB = 103.15"
	elif "новости" in text.lower() or "news" in text.lower():
		tool_result = "MCP инструмент 'news': Последние новости: ИИ развивается быстрыми темпами, новые технологии в разработке"
	elif "поиск" in text.lower() or "search" in text.lower() or "найди" in text.lower():
		tool_result = f"MCP инструмент 'search': Найдено 1,234,567 результатов по запросу '{text}'"
	else:
		tool_result = f"MCP инструмент 'generic': Обработан запрос '{text}' через MCP протокол"
	
	print(f"🌐 Результат MCP инструмента: {tool_result}")
	return {"tool_result": tool_result}


def _respond_node(state: AgentState) -> Dict[str, Any]:
	"""Создает финальный ответ, опционально используя результат инструмента.

	Этот узел использует GigaChat для генерации ответов.
	"""
	tool_result = state.get("tool_result")
	last_user = None
	for msg in reversed(state.get("messages", [])):
		if isinstance(msg, HumanMessage):
			last_user = msg.content
			break

	if tool_result:
		# Если у нас есть результат инструмента, включаем его в ответ
		print(f"📊 Результат инструмента: {tool_result}")
		prompt = f"""Пользователь спросил: "{last_user}"
Результат инструмента: {tool_result}

Дай краткий и дружелюбный ответ, включающий результат инструмента."""
		response = giga.chat(prompt)
		reply = response.choices[0].message.content
		print("💬 Используется GigaChat с результатом инструмента")
	else:
		# Нет результата инструмента - общий разговор
		prompt = f"""Ты полезный AI ассистент. Пользователь написал: "{last_user}"
Дай краткий и полезный ответ."""
		response = giga.chat(prompt)
		reply = response.choices[0].message.content
		print("💬 Используется GigaChat для общего ответа")

	return {"messages": [AIMessage(content=reply)], "tool_result": None}


# Строим граф
_graph_builder = StateGraph(AgentState)
_graph_builder.add_node("input", _passthrough_node)
_graph_builder.add_node("tool", _math_tool_node)
_graph_builder.add_node("mcp", _mcp_tool_node)
_graph_builder.add_node("respond", _respond_node)

_graph_builder.add_conditional_edges(
	"input",
	_decide_next,
	{
		"tool": "tool",
		"mcp": "mcp", 
		"respond": "respond",
	},
)

_graph_builder.add_edge("tool", "respond")
_graph_builder.add_edge("mcp", "respond")
_graph_builder.add_edge("respond", END)
_graph_builder.set_entry_point("input")

AGENT = _graph_builder.compile()


def run_agent(user_text: str) -> str:
	"""Запускает минимального агента и возвращает текст ответа ассистента.

	Параметры:
		user_text: сообщение пользователя для обработки
	"""
	state: AgentState = {
		"messages": [HumanMessage(content=user_text)],
		"tool_result": None,
		"mcp_tools": None,
	}
	out = AGENT.invoke(state)
	messages = out.get("messages", [])
	last = messages[-1] if messages else AIMessage(content="")
	return str(getattr(last, "content", ""))


if __name__ == "__main__":
	# Быстрая ручная проверка: `python -m tests.tools`
	import sys

	text = " ".join(sys.argv[1:]) or "сложи 2 и 3 и 4"
	print(run_agent(text))