# === system_prompt.j2 (Jinja2 шаблон) ===
Ты — AI-агент «Аналитик» для платформы qantcore.space.
[СЛОЙ 1: РОЛЬ]
Твоя роль: data-analyst. Ты анализируешь данные, строишь
гипотезы и выдаёшь рекомендации на основе цифр.

[СЛОЙ 2: КОНТЕКСТ]
Текущая дата: {{ current_date }}
Пользователь: {{ user_name }}, роль: {{ user_role }}
База данных: PostgreSQL с таблицами orders, users, products

[СЛОЙ 3: ПРАВИЛА]
1. Всегда проверяй результат вычислений (sanity check)
2. Если данных недостаточно — запроси уточнение, НЕ угадывай
3. Отвечай на русском языке, числа форматируй с разделителями
4. В ответе указывай источник данных и метод расчёта
5. При ошибке в данных — сообщи об этом и предложи обходной путь

[СЛОЙ 4: ИНСТРУМЕНТЫ]
Доступные функции:
- query_db(sql: str) → dict  # SQL SELECT к PostgreSQL
- calculate(expr: str) → float  # Математические вычисления
- plot_chart(data: dict) → url  # Генерация графика
- web_search(query: str) → list  # Поиск в интернете

[СЛОЙ 5: ФОРМАТ ОТВЕТА]
Формат:
{
  "analysis": "<краткий анализ>",
  "data": "<ключевые цифры>",
  "recommendation": "<рекомендация>",
  "confidence": 0.0-1.0
}

# === cot_agent.py ===
import re
from openai import AsyncOpenAI

# CoT-обёртка: шаги → ответ
COT_PROMPT = """
Решай задачу пошагово. Для каждого шага используй формат:

ШАГ {N}: <описание шага>
МЫСЛЬ: <твои рассуждения>
ДЕЙСТВИЕ: <вызов функции или вычисление>
РЕЗУЛЬТАТ: <что получилось>

После всех шагов:

ИТОГОВЫЙ ОТВЕТ: <финальный ответ>

Задача: {task}
"""

class COTAgent:
    def __init__(self, client: AsyncOpenAI, model="gpt-4o"):
        self.client = client
        self.model = model

    async def solve(self, task: str) -> dict:
        prompt = COT_PROMPT.format(task=task)
        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3
        )
        text = response.choices[0].message.content

        # Парсим шаги и итоговый ответ
        steps = re.findall(
            r'ШАГ (\d+):(.+?)(?=ШАГ \d+:|ИТОГОВЫЙ ОТВЕТ:|$)',
            text, re.DOTALL
        )
        final = re.search(
            r'ИТОГОВЫЙ ОТВЕТ:\s*(.+)', text, re.DOTALL
        )

        return {
            "steps": [(int(n), s.strip()) for n, s in steps],
            "answer": final.group(1) if final else "не найдено"
        }

# === react_agent.py ===
import json
from typing import Any

REACT_SYSTEM_PROMPT = """
Ты агент, работающий по циклу ReAct. На каждом шаге выбирай:

THOUGHT: твои рассуждения о том, что нужно сделать
ACTION: вызов инструмента (в JSON):
  {"tool": "имя_функции", "args": {...}}
STOP: когда ответ готов, верни:
  {"action": "finish", "answer": "твой ответ"}

Инструменты:
- search(query: str) → список результатов
- calculate(expression: str) → число
- get_weather(city: str) → погода
"""

class ReActAgent:
    def __init__(self, llm, tools: dict):
        self.llm = llm
        self.tools = tools
        self.max_steps = 10

    async def run(self, task: str) -> str:
        messages = [
            {"role": "system", "content": REACT_SYSTEM_PROMPT},
            {"role": "user", "content": f"Задача: {task}"}
        ]

        for step in range(self.max_steps):
            resp = await self.llm.chat(messages)
            action = self._parse_action(resp)

            if action["action"] == "finish":
                return action["answer"]

            # Выполняем инструмент и добавляем Observation
            tool_name = action["tool"]
            tool_args = action["args"]
            result = await self.tools[tool_name](**tool_args)
            messages.append({
                "role": "assistant",
                "content": f"OBSERVATION: {json.dumps(result, ensure_ascii=False)}"
            })

        return "Достигнут лимит шагов"

# === tot_agent.py ===
import asyncio

TOT_THINK_PROMPT = """
Задача: {task}
Текущий контекст: {context}

Сгенерируй ровно 3 различных варианта следующего шага.
Для каждого варианта объясни почему он может быть хорош.

Формат:
ВАРИАНТ 1: <описание>
ОЦЕНКА 1: <почему это хороший вариант>

ВАРИАНТ 2: <описание>
ОЦЕНКА 2: <почему это хороший вариант>

ВАРИАНТ 3: <описание>
ОЦЕНКА 3: <почему это хороший вариант>

Теперь выбери ЛУЧШИЙ вариант и напиши:
ВЫБОР: <номер лучшего варианта>
ПРИЧИНА: <почему выбрал именно его>
"""

class TreeOfThought:
    def __init__(self, llm, max_depth=3, branches=3):
        self.llm = llm
        self.max_depth = max_depth
        self.branches = branches

    async def solve(self, task: str) -> dict:
        return await self._search(task, context="", depth=0)

    async def _search(self, task, context, depth):
        if depth >= self.max_depth:
            return await self._evaluate_final(task, context)

        prompt = TOT_THINK_PROMPT.format(task=task, context=context)
        response = await self.llm.generate(prompt)

        # Парсим варианты и выбор
        variants = self._parse_variants(response)
        best_idx = self._parse_choice(response)

        # Рекурсивно углубляем ЛУЧШИЙ вариант
        best_variant = variants[best_idx]
        new_context = f"{context}\n→ {best_variant}"
        return await self._search(task, new_context, depth + 1)

# === few_shot_store.py ===
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity

class DynamicFewShot:
    """Хранилище примеров с семантическим поиском."""
    def __init__(self, embedder):
        self.embedder = embedder
        self.examples = []        # list of dict: {q, a, embedding}

    def add(self, question: str, answer: str):
        emb = self.embedder.encode(question)
        self.examples.append({"q": question, "a": answer, "emb": emb})

    def retrieve(self, query: str, k=3) -> list:
        if not self.examples:
            return []
        q_emb = self.embedder.encode(query)
        sims = cosine_similarity([q_emb], [e["emb"] for e in self.examples])[0]
        top_k = np.argsort(sims)[-k:][::-1]
        return [self.examples[i] for i in top_k]

    def format_prompt(self, query: str, k=3) -> str:
        examples = self.retrieve(query, k)
        parts = ["Примеры:\n"]
        for i, ex in enumerate(examples, 1):
            parts.append(f"{i}. Вопрос: {ex['q']}\n   Ответ: {ex['a']}")
        parts.append(f"\nТеперь ответь на:\n{query}")
        return "\n".join(parts)

# === eval_prompts.py ===
import json, asyncio, time
from sentence_transformers import SentenceTransformer

class PromptEvaluator:
    def __init__(self, agent, testset_path: str):
        self.agent = agent
        self.tests = self._load_tests(testset_path)
        self.model = SentenceTransformer("all-MiniLM-L6-v2")

    async def evaluate(self) -> dict:
        results = {"total": len(self.tests), "passed": 0, "semantic_score": 0.0, "latency": []}
        for test in self.tests:
            start = time.monotonic()
            actual = await self.agent.run(test["question"])
            elapsed = time.monotonic() - start

            # Семантическая близость ответа к эталону
            emb_actual = self.model.encode(actual)
            emb_expected = self.model.encode(test["expected"])
            score = cosine_similarity([emb_actual], [emb_expected])[0][0]

            results["semantic_score"] += score
            results["latency"].append(elapsed)
            if score > 0.8:
                results["passed"] += 1

        results["semantic_score"] /= results["total"]
        results["avg_latency"] = sum(results["latency"]) / len(results["latency"])
        return results

# Пример testset.jsonl:
# {"question": "2+2*2", "expected": "6 (умножение выполняется первым)"}
# {"question": "Столица Японии", "expected": "Токио"}