core_agent.py

"""
core_agent.py — TEKDEV Bot Core
Handles all text generation and agentic logic via Gemma3:1B (InferenceClient).
"""

from huggingface_hub import InferenceClient
from config import HF_TOKEN, MODEL_ID, MAX_TOKENS, TEMPERATURE
from personality import get_system_prompt
from web_search import search_web

# ── Client ─────────────────────────────────────────────────────────────────────
client = InferenceClient(token=HF_TOKEN)

# ── Search trigger keywords ────────────────────────────────────────────────────
SEARCH_TRIGGERS = [
    "search", "find", "look up", "latest", "news", "current",
    "today", "price", "weather", "who is", "what is", "how much",
    "when did", "recent", "update", "right now",
]

def _should_search(message: str) -> bool:
    lowered = message.lower()
    return any(kw in lowered for kw in SEARCH_TRIGGERS)


def _build_messages(user_message: str, history: list[dict], web_context: str | None) -> list[dict]:
    """Assemble the full message array for the model."""
    system_prompt = get_system_prompt()
    messages = [{"role": "system", "content": system_prompt}]

    # Inject conversation history (trimmed to last 10 turns to stay within context)
    messages += history[-10:]

    if web_context:
        # Prepend search results as a tool note so the model stays grounded
        enriched = (
            f"[WEB SEARCH RESULTS]\n{web_context}\n\n"
            f"Use the above information to answer:\n{user_message}"
        )
        messages.append({"role": "user", "content": enriched})
    else:
        messages.append({"role": "user", "content": user_message})

    return messages


def run_agent(user_message: str, history: list[dict] | None = None) -> str:
    """
    Main entry point.  Returns the assistant reply as a plain string.

    Args:
        user_message: The latest user input.
        history:      List of {"role": "user"|"assistant", "content": "..."} dicts.
    """
    if history is None:
        history = []

    web_context: str | None = None
    if _should_search(user_message):
        try:
            web_context = search_web(user_message)
        except Exception as exc:
            web_context = f"(Search unavailable: {exc})"

    messages = _build_messages(user_message, history, web_context)

    try:
        response = client.chat_completion(
            model=MODEL_ID,
            messages=messages,
            max_tokens=MAX_TOKENS,
            temperature=TEMPERATURE,
        )
        reply = response.choices[0].message.content.strip()
    except Exception as exc:
        reply = f"⚠️ Model error: {exc}"

    return reply


def stream_agent(user_message: str, history: list[dict] | None = None):
    """
    Streaming variant — yields text chunks for real-time output.
    Use this if you want to stream responses to Telegram or a UI.
    """
    if history is None:
        history = []

    web_context: str | None = None
    if _should_search(user_message):
        try:
            web_context = search_web(user_message)
        except Exception:
            pass

    messages = _build_messages(user_message, history, web_context)

    for chunk in client.chat_completion(
        model=MODEL_ID,
        messages=messages,
        max_tokens=MAX_TOKENS,
        temperature=TEMPERATURE,
        stream=True,
    ):
        delta = chunk.choices[0].delta.content
        if delta:
            yield delta