agent.py

"""
eraCore.agent — AI-агент с LLM, поиском, памятью и файлами.

Пример::

    import asyncio
    from eraCore import Agent
    from eraCore.config import AgentConfig

    # Без LLM (поиск + кэш)
    agent = Agent()

    # С Ollama
    agent = Agent(config=AgentConfig(
        llm_url="http://localhost:11434/v1",
        llm_model="qwen2.5:3b",
    ))

    # С локальным GGUF
    agent = Agent(config=AgentConfig(llm_path="./models/qwen-3b.gguf"))

    async def main():
        answer = await agent.ask("Что такое нейронная сеть?")
        print(answer)

        async for chunk in await agent.ask("Объясни backprop", stream=True):
            print(chunk, end="", flush=True)

        answer = await agent.ask_file("data.csv", "Что за данные?")

    asyncio.run(main())
"""

from __future__ import annotations

import asyncio
from typing import AsyncGenerator, List, Optional, Union

from eraCore.config import AgentConfig
from eraCore.tools.search import SearchTool
from eraCore.tools.brain import BrainTool
from eraCore.tools.file import FileTool


class Agent:
    """
    AI-агент. Три инструмента + LLM (опционально).

    Без LLM:  search → агрегирует результаты как текст.
    С LLM:    search → контекст → LLM → ответ.
    """

    def __init__(
        self,
        config: Optional[AgentConfig] = None,
        # legacy-параметры для совместимости
        *,
        knowledge_path: str = ".era_knowledge.json",
        searx_url: str = "https://searx.be/search",
        cache_threshold: float = 0.82,
        context_limit: int = 20,
    ) -> None:
        # Конфиг: если передан — берём его, иначе собираем из legacy-параметров
        if config is None:
            config = AgentConfig(
                knowledge_path=knowledge_path,
                searx_url=searx_url,
                cache_threshold=cache_threshold,
                context_limit=context_limit,
            )
        self.config = config

        # Инструменты
        self.search = SearchTool(
            searx_url=config.searx_url,
            timeout=config.search_timeout,
            max_results=config.search_max_results,
        )
        self.brain = BrainTool(
            knowledge_path=config.knowledge_path,
            threshold=config.cache_threshold,
            context_limit=config.context_limit,
        )
        self.file = FileTool()

        # LLM (ленивая инициализация)
        self._llm = None
        if config.has_llm():
            self._init_llm()

        # Режимы
        self._modes: dict = {
            "original": self._mode_original,
            "coder":    self._mode_coder,
            "expert":   self._mode_expert,
        }

    # ──────────────────── LLM ────────────────────

    def _init_llm(self) -> None:
        from eraCore.llm import LLM
        cfg = self.config
        self._llm = LLM(
            url=cfg.llm_url,
            model=cfg.llm_model,
            path=cfg.llm_path,
            api_key=cfg.llm_api_key,
            max_tokens=cfg.llm_max_tokens,
            temperature=cfg.llm_temperature,
            system=cfg.llm_system,
        )
        if cfg.debug:
            print(f"[Agent] LLM: {self._llm.info()}")

    def set_llm(
        self,
        url: Optional[str] = None,
        model: str = "qwen2.5:3b",
        path: Optional[str] = None,
        **kwargs,
    ) -> None:
        """Подключить LLM в любой момент после создания агента."""
        from eraCore.llm import LLM
        self._llm = LLM(url=url, model=model, path=path, **kwargs)
        self.config.llm_url = url
        self.config.llm_model = model
        self.config.llm_path = path

    # ──────────────────── ASK ────────────────────

    async def ask(
        self,
        question: str,
        mode: str = "original",
        stream: bool = False,
    ) -> Union[str, AsyncGenerator[str, None]]:
        """
        Задать вопрос агенту.

        Flow:
          1. brain.recall()         → кэш-хит → сразу ответ
          2. режим(question)        → search → [LLM] → ответ
          3. brain.remember/push()  → сохраняем результат
        """
        question = (question or "").strip()
        if not question:
            txt = "Задайте вопрос."
            return txt if not stream else self._as_stream(txt)

        # 1. Кэш
        cached = self.brain.recall(question)
        if cached:
            if self.config.debug:
                print(f"[Agent] cache hit: {question[:50]}")
            self.brain.push("assistant", cached)
            return cached if not stream else self._as_stream(f"[из памяти]\n{cached}")

        # 2. Режим
        handler = self._modes.get(mode, self._modes["original"])
        result = await handler(question, stream=stream)

        # 3. Сохраняем (только для не-стримов — стримы сохраняют сами)
        if not stream:
            self.brain.remember(question, result)
            self.brain.push("user", question)
            self.brain.push("assistant", result)

        return result

    async def ask_file(
        self,
        path: str,
        instruction: str = "Опиши содержимое файла",
    ) -> str:
        """Прочитать файл и ответить на вопрос о нём."""
        try:
            content = self.file.read(path)
        except FileNotFoundError:
            return f"Файл не найден: {path}"

        info = self.file.info(path)
        preview = content[:3000]

        if self._llm:
            prompt = (
                f"Файл: {info['name']} ({info['size_kb']} KB, тип: {info['type']})\n\n"
                f"Содержимое:\n{preview}\n\n"
                f"Задача: {instruction}"
            )
            return await self._llm.agenerate(prompt)

        # Без LLM — структурированный превью
        return (
            f"**{info['name']}** | {info['size_kb']} KB | {info['type']}"
            f" | {info.get('lines', '?')} строк\n\n"
            f"**Задача:** {instruction}\n\n"
            f"```\n{preview[:1000]}\n```"
        )

    # ──────────────────── РЕЖИМЫ ────────────────────

    async def _mode_original(self, question: str, stream: bool = False) -> Union[str, AsyncGenerator]:
        """Поиск → [LLM с контекстом] → ответ."""
        results = await asyncio.to_thread(self.search.run, question)
        ctx = self._build_context(results, 2000)
        history = self.brain.context_text(last_n=6)

        if self._llm:
            prompt = self._build_prompt(question, ctx, history)
            if stream:
                return self._llm_stream_with_save(question, prompt, self._fallback_original(results))
            answer = await self._llm.agenerate(prompt)
            if self.config.fallback_on_llm_error and self._is_llm_error(answer):
                return self._fallback_original(results)
            return answer

        return self._fallback_original(results)

    def _fallback_original(self, results: list) -> str:
        """Результат без LLM для original-режима."""
        if not results:
            return "По вашему запросу ничего не найдено."
        lines = ["**Результаты поиска:**\n"]
        for i, r in enumerate(results[:4], 1):
            title = r.get("title", "")
            snippet = (r.get("content") or r.get("snippet") or "")[:300]
            url = r.get("url", "")
            lines.append(f"{i}. **{title}**")
            if snippet:
                lines.append(snippet)
            if url:
                lines.append(f"   → {url}")
            lines.append("")
        return "\n".join(lines)

    async def _mode_coder(self, question: str, stream: bool = False) -> Union[str, AsyncGenerator]:
        """Поиск code-примеров → [LLM] → код."""
        results = await asyncio.to_thread(
            self.search.multi, [question, f"{question} python example code"]
        )
        ctx = self._build_context(results, 2000)
        history = self.brain.context_text(last_n=4)

        if self._llm:
            prompt = self._build_prompt(
                question, ctx, history,
                system="Ты эксперт-программист. Давай точные, рабочие примеры кода с пояснениями."
            )
            if stream:
                return self._llm_stream_with_save(question, prompt, self._fallback_coder(question, ctx))
            answer = await self._llm.agenerate(prompt)
            if self.config.fallback_on_llm_error and self._is_llm_error(answer):
                return self._fallback_coder(question, ctx)
            return answer

        return self._fallback_coder(question, ctx)

    def _fallback_coder(self, question: str, ctx: str) -> str:
        """Результат без LLM для coder-режима."""
        return (
            f"**[CODER]** {question}\n\n"
            f"Контекст из поиска:\n```\n{ctx[:1000]}\n```\n\n"
            "_Подключи LLM: `agent.set_llm(url='http://localhost:11434/v1')`_"
        )

    async def _mode_expert(self, question: str, stream: bool = False) -> Union[str, AsyncGenerator]:
        """Глубокий поиск → [LLM] → подробный ответ + источники."""
        queries = [question, f"{question} подробно", f"{question} примеры"]
        results = await asyncio.to_thread(self.search.multi, queries)
        ctx = self._build_context(results, 3000)
        sources = self._build_sources(results)
        history = self.brain.context_text(last_n=6)

        if self._llm:
            prompt = self._build_prompt(
                question, ctx, history,
                system="Ты глубокий эксперт. Давай структурированные, подробные ответы с примерами."
            )
            if stream:
                return self._expert_stream_safe(question, prompt, ctx, sources)
            answer = await self._llm.agenerate(prompt)
            if self.config.fallback_on_llm_error and self._is_llm_error(answer):
                return self._fallback_expert(question, ctx, sources)
            return f"{answer}\n\n{sources}" if sources else answer

        return self._fallback_expert(question, ctx, sources)

    def _fallback_expert(self, question: str, ctx: str, sources: str) -> str:
        """Результат без LLM для expert-режима."""
        return f"**[EXPERT]** {question}\n\n{ctx}\n\n{sources}"

    # ──────────────────── PROMPT BUILDER ────────────────────

    def _build_prompt(
        self,
        question: str,
        context: str,
        history: str,
        system: Optional[str] = None,
    ) -> str:
        parts = []
        if system:
            parts.append(f"[Инструкция]\n{system}\n")
        if history:
            parts.append(f"[История разговора]\n{history}\n")
        if context:
            parts.append(f"[Контекст из поиска]\n{context}\n")
        parts.append(f"[Вопрос]\n{question}")
        return "\n\n".join(parts)

    async def _llm_stream_with_save(
        self, question: str, prompt: str, fallback: str
    ) -> AsyncGenerator[str, None]:
        """Стриминг через LLM + сохранение в brain. Fallback при ошибке первого чанка."""
        gen = self._llm.astream(prompt)
        try:
            first = await gen.__anext__()
        except StopAsyncIteration:
            yield fallback
            return

        if self.config.fallback_on_llm_error and self._is_llm_error(first):
            yield fallback
            return

        acc = first
        yield first
        async for chunk in gen:
            acc += chunk
            yield chunk
        self.brain.remember(question, acc)
        self.brain.push("user", question)
        self.brain.push("assistant", acc)

    async def _expert_stream_safe(
        self, question: str, prompt: str, ctx: str, sources: str
    ) -> AsyncGenerator[str, None]:
        """Стриминг expert с fallback при ошибке."""
        gen = self._llm.astream(prompt)
        try:
            first = await gen.__anext__()
        except StopAsyncIteration:
            yield self._fallback_expert(question, ctx, sources)
            return

        if self.config.fallback_on_llm_error and self._is_llm_error(first):
            yield self._fallback_expert(question, ctx, sources)
            return

        yield first
        async for chunk in gen:
            yield chunk
        if sources:
            yield f"\n\n{sources}"

    # ──────────────────── ADD MODE ────────────────────

    def add_mode(self, name: str, fn) -> None:
        """
        Зарегистрировать свой режим.

        fn может быть:
          async def my_mode(question: str, stream: bool = False) -> str | AsyncGenerator
        """
        self._modes[name] = fn

    # ──────────────────── INFO ────────────────────

    def modes(self) -> List[str]:
        return list(self._modes.keys())

    def stats(self) -> dict:
        return {
            "modes": self.modes(),
            "llm": self._llm.info() if self._llm else None,
            "brain": self.brain.stats(),
            "config": {
                "llm_url": self.config.llm_url,
                "llm_model": self.config.llm_model,
                "has_llm": self.config.has_llm(),
            },
        }

    # ──────────────────── STATIC HELPERS ────────────────────

    @staticmethod
    def _build_context(results: list, max_chars: int) -> str:
        parts, total = [], 0
        for r in results:
            text = (r.get("content") or r.get("snippet") or "").strip()
            if text and total < max_chars:
                chunk = text[: max_chars - total]
                parts.append(chunk)
                total += len(chunk)
        return "\n\n".join(parts) or ""

    @staticmethod
    def _is_llm_error(text: str) -> bool:
        """Проверить, является ли ответ сообщением об ошибке LLM."""
        return isinstance(text, str) and text.startswith("[Ошибка]")

    @staticmethod
    def _build_sources(results: list) -> str:
        seen, sources = set(), []
        for r in results:
            url = r.get("url", "")
            if url and url not in seen:
                seen.add(url)
                sources.append(f"- [{r.get('title', 'Источник')[:80]}]({url})")
        return ("**Источники:**\n" + "\n".join(sources[:5])) if sources else ""

    @staticmethod
    async def _as_stream(text: str) -> AsyncGenerator[str, None]:
        for ch in text:
            yield ch
            await asyncio.sleep(0.004)