app.py

import asyncio
import json
import os
import time
import uvicorn
import argparse
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional

import httpx
from fastapi import FastAPI, Request, Response
from fastapi.responses import StreamingResponse
from pydantic import BaseModel, Field

from research import ResearchOrchestrator, build_context_block
from llm_client import stream_upstream_sse

from dotenv import load_dotenv
from zoneinfo import ZoneInfo

# NEW: multilingual trigger detection
from lang_signals import recency_signals, howto_signals

load_dotenv()

APP_NAME = "web_relay"
app = FastAPI(title="Web Relay", version="1.0.0")

UPSTREAM_TYPE = os.getenv("UPSTREAM_TYPE", "llama").lower()        # llama | lmstudio
UPSTREAM_URL  = os.getenv("UPSTREAM_URL", "http://127.0.0.1:8080")
DEFAULT_MODEL = os.getenv("DEFAULT_MODEL", "gemma-3-12b-it-ud@q8_k_xl")
CTX_BUDGET    = int(os.getenv("CONTEXT_BUDGET_CHARS", "7000"))
LOCAL_TZ_NAME = os.getenv("LOCAL_TZ", "Europe/Zurich")

REQUEST_TIMEOUT = float(os.getenv("REQUEST_TIMEOUT", "20"))


def now_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def build_time_context() -> Dict[str, str]:
    """
    Provides a consistent time anchor:
    - utc_iso: 2025-11-18T09:10:11Z
    - local_iso: 2025-11-18T10:10:11+01:00
    - local_date: 2025-11-18
    - tz_name: Europe/Zurich (or fallback UTC)
    """
    tz_name = LOCAL_TZ_NAME
    try:
        tz = ZoneInfo(tz_name)
    except Exception:
        tz = timezone.utc
        tz_name = "UTC"

    utc_iso = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")
    local_now = datetime.now(tz)
    local_iso = local_now.isoformat()
    local_date = local_now.strftime("%Y-%m-%d")
    return {
        "utc_iso": utc_iso,
        "local_iso": local_iso,
        "local_date": local_date,
        "tz_name": tz_name
    }


class ChatMessage(BaseModel):
    role: str
    content: str


class RelayBody(BaseModel):
    model: str = Field(default_factory=lambda: DEFAULT_MODEL)
    messages: List[ChatMessage]
    temperature: Optional[float] = 0.7
    top_p: Optional[float] = 0.95
    top_k: Optional[int] = 60
    presence_penalty: Optional[float] = 0.0
    frequency_penalty: Optional[float] = 0.0
    stream: Optional[bool] = True
    # Optional budget, if the client sends it explicitly
    char_budget: Optional[int] = None


@app.get("/health")
async def health():
    tctx = build_time_context()
    return {
        "status": "ok",
        "ts_utc": tctx["utc_iso"],
        "today_local": tctx["local_date"],
        "tz": tctx["tz_name"],
        "upstream": UPSTREAM_TYPE,
        "url": UPSTREAM_URL
    }


def need_web(query: str) -> bool:
    """
    Multilingual heuristic to decide whether web access is useful.
    Uses language-driven recency/how-to signals (30 languages) + a few universal extras.
    """
    q = (query or "").lower()
    # Multilingual cues
    if recency_signals(q) or howto_signals(q):
        return True

    # Universal extras
    if "http://" in q or "https://" in q:
        return True

    # Optional generic shopping/review cues (kept in English only to avoid false positives)
    extras = [
        "comparison", "compare", "vs.", "best", "top", "review", "reviews", "test",
        "download link", "install guide", "installation guide", "api reference"
    ]
    return any(t in q for t in extras)


def strip_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
    """If the client sends 'tool' lines, we turn them into neutral hints."""
    out: List[Dict[str, Any]] = []
    for m in messages:
        role = m.get("role", "")
        if role == "tool":
            # Instrument tool content as additional information for the LLM
            out.append({"role": "user", "content": f"(Note from tool) {m.get('content', '')}"})
        else:
            out.append(m)
    return out


def build_system_web_guidance_with_time(tctx: Dict[str, str]) -> str:
    """
    Dynamic guidance with an explicit date/time anchor for web searches.
    This guidance is ONLY added when need_web(...) is True.
    """
    base = (
        "You are a careful assistant. When the context block `<<<CONTEXT>>>` "
        "contains web sources, use them to answer the question. "
        "Cite evidence inline with square brackets like [1], [2], matching the sources in the context. "
        "If information is unclear or sources conflict, explain this openly and summarize the situation. "
        "If a step-by-step guide is requested, provide a clear, numbered set of steps based on the sources. "
        "Do not invent quotes or references."
    )
    time_anchor = (
        "\n\nTIME ANCHOR:\n"
        f"- Current date (local, {tctx['tz_name']}): {tctx['local_date']}\n"
        f"- Current time (local, {tctx['tz_name']}): {tctx['local_iso']}\n"
        f"- Current time (UTC): {tctx['utc_iso']}\n"
        "IMPORTANT: Treat expressions like 'today', 'currently', or 'now' strictly relative to this date/time anchor. "
        "If the user explicitly mentions another date, then relative expressions refer to that given date instead."
    )
    return base + time_anchor


# Legacy variant (if no web search is needed)
SYSTEM_WEB_GUIDANCE = (
    "You are a careful assistant. When the context block `<<<CONTEXT>>>` "
    "contains web sources, use them to answer the question. "
    "Cite evidence inline with square brackets like [1], [2], matching the sources in the context. "
    "If information is unclear or sources conflict, explain this openly and summarize the situation. "
    "If a step-by-step guide is requested, provide a clear, numbered set of steps based on the sources. "
    "Do not invent quotes or references."
)


@app.post("/relay")
async def relay(req: Request):
    body = await req.json()
    data = RelayBody(**body)

    # 1) extract last user message
    raw_msgs = [m.model_dump() if isinstance(m, ChatMessage) else m for m in data.messages]
    msgs = strip_tool_messages(raw_msgs)
    last_user_idx = max(i for i, m in enumerate(msgs) if m.get("role") == "user") if any(
        m.get("role") == "user" for m in msgs
    ) else -1
    last_user_text = msgs[last_user_idx]["content"] if last_user_idx >= 0 else ""

    # Precompute time anchor (may be used both in system prompt and context)
    tctx = build_time_context()

    # 2) Web plan
    char_budget = int(data.char_budget) if data.char_budget else CTX_BUDGET
    context_block = ""
    will_use_web = bool(last_user_text.strip()) and need_web(last_user_text)

    if will_use_web:
        researcher = ResearchOrchestrator()
        try:
            digest = await researcher.research_and_digest(last_user_text, max_chars=char_budget)
            if digest and digest["items"]:
                context_block = build_context_block(
                    query=last_user_text,
                    items=digest["items"],
                    generated_at=tctx["utc_iso"],
                    current_local_iso=tctx["local_iso"],
                    current_local_date=tctx["local_date"],
                    tz_name=tctx["tz_name"],
                    lang_hint="en",
                    budget_chars=char_budget,
                )
        except Exception as e:
            # Context stays empty; we still answer – upstream can respond generically.
            context_block = f"<<<CONTEXT>>>\n(Note: Web search failed: {e})\n<<<END_CONTEXT>>>"

    # 3) Rebuild messages for upstream
    new_messages: List[Dict[str, str]] = []
    # Keep original system prompt (if any) at the top + add guidance
    existing_sys = next((m for m in msgs if m.get("role") == "system"), None)
    if existing_sys:
        new_messages.append(existing_sys)
        # For web searches always use the dated guidance, otherwise the legacy variant
        sys_text = build_system_web_guidance_with_time(tctx) if will_use_web else SYSTEM_WEB_GUIDANCE
        new_messages.append({"role": "system", "content": sys_text})
    else:
        sys_text = build_system_web_guidance_with_time(tctx) if will_use_web else SYSTEM_WEB_GUIDANCE
        new_messages.append({"role": "system", "content": sys_text})

    for i, m in enumerate(msgs):
        if i == last_user_idx and context_block:
            merged = f"{context_block}\n\n{m['content']}"
            new_messages.append({"role": "user", "content": merged})
        elif m.get("role") in ("user", "assistant"):
            new_messages.append(m)

    # 4) Stream upstream and forward SSE 1:1 to client
    async def sse_iter():
        try:
            async for sse_line in stream_upstream_sse(
                upstream_type=UPSTREAM_TYPE,
                upstream_url=UPSTREAM_URL,
                model=data.model or DEFAULT_MODEL,
                messages=new_messages,
                temperature=data.temperature,
                top_p=data.top_p,
                top_k=data.top_k,
                presence_penalty=data.presence_penalty,
                frequency_penalty=data.frequency_penalty,
            ):
                yield sse_line
        except httpx.HTTPError as e:
            # Emit error message as a single JSON chunk delta
            msg = f"Web relay error: {e}"
            payload = {
                "id": "relay-error",
                "object": "chat.completion.chunk",
                "model": data.model,
                "choices": [
                    {"index": 0, "delta": {"content": f"\n⚠️ {msg}"}, "finish_reason": None}
                ],
            }
            yield f"data: {json.dumps(payload, ensure_ascii=False)}\n\n"
        finally:
            # Completion marker
            yield "data: [DONE]\n\n"

    headers = {
        "Content-Type": "text/event-stream; charset=utf-8",
        "Cache-Control": "no-cache, no-store, must-revalidate",
        "Connection": "keep-alive",
    }
    return StreamingResponse(sse_iter(), headers=headers, media_type="text/event-stream")


@app.post("/relay_once")
async def relay_once(req: Request):
    """
    Variant for classic clients (e.g. C# frontend).
    No SSE streams – returns a single JSON response after completion.
    """
    body = await req.json()
    data = RelayBody(**body)

    # 1) extract last user message
    raw_msgs = [m.model_dump() if isinstance(m, ChatMessage) else m for m in data.messages]
    msgs = strip_tool_messages(raw_msgs)
    last_user_idx = max(i for i, m in enumerate(msgs) if m.get("role") == "user") if any(
        m.get("role") == "user" for m in msgs
    ) else -1
    last_user_text = msgs[last_user_idx]["content"] if last_user_idx >= 0 else ""

    # Time anchor (same as in /relay)
    tctx = build_time_context()

    # 2) Web plan
    char_budget = int(data.char_budget) if data.char_budget else CTX_BUDGET
    context_block = ""
    will_use_web = bool(last_user_text.strip()) and need_web(last_user_text)

    if will_use_web:
        researcher = ResearchOrchestrator()
        try:
            digest = await researcher.research_and_digest(last_user_text, max_chars=char_budget)
            if digest and digest["items"]:
                context_block = build_context_block(
                    query=last_user_text,
                    items=digest["items"],
                    generated_at=tctx["utc_iso"],
                    current_local_iso=tctx["local_iso"],
                    current_local_date=tctx["local_date"],
                    tz_name=tctx["tz_name"],
                    lang_hint="en",
                    budget_chars=char_budget,
                )
        except Exception as e:
            context_block = f"<<<CONTEXT>>>\n(Note: Web search failed: {e})\n<<<END_CONTEXT>>>"

    # 3) Rebuild messages
    new_messages: List[Dict[str, str]] = []
    existing_sys = next((m for m in msgs if m.get("role") == "system"), None)
    if existing_sys:
        new_messages.append(existing_sys)
        sys_text = build_system_web_guidance_with_time(tctx) if will_use_web else SYSTEM_WEB_GUIDANCE
        new_messages.append({"role": "system", "content": sys_text})
    else:
        sys_text = build_system_web_guidance_with_time(tctx) if will_use_web else SYSTEM_WEB_GUIDANCE
        new_messages.append({"role": "system", "content": sys_text})

    for i, m in enumerate(msgs):
        if i == last_user_idx and context_block:
            merged = f"{context_block}\n\n{m['content']}"
            new_messages.append({"role": "user", "content": merged})
        elif m.get("role") in ("user", "assistant"):
            new_messages.append(m)

    # 4) Upstream request (non-streaming)
    try:
        async with httpx.AsyncClient(timeout=httpx.Timeout(REQUEST_TIMEOUT, connect=REQUEST_TIMEOUT)) as client:
            response = await client.post(
                f"{UPSTREAM_URL}/v1/chat/completions",
                json={
                    "model": data.model or DEFAULT_MODEL,
                    "messages": new_messages,
                    "temperature": data.temperature,
                    "top_p": data.top_p,
                    "top_k": data.top_k,
                    "presence_penalty": data.presence_penalty,
                    "frequency_penalty": data.frequency_penalty,
                    "stream": False,
                },
            )
            response.raise_for_status()
            upstream_json = response.json()
    except Exception as e:
        return {
            "choices": [
                {
                    "message": {
                        "role": "assistant",
                        "content": f"⚠️ Error in web relay: {str(e)}",
                    }
                }
            ]
        }

    # 5) Extract answer and return it in a C#-compatible format
    try:
        content = upstream_json.get("choices", [{}])[0].get("message", {}).get("content", "")
    except Exception:
        content = str(upstream_json)

    return {
        "choices": [
            {
                "message": {
                    "role": "assistant",
                    "content": content,
                }
            }
        ]
    }


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Start the Aurelia Web Relay.")
    parser.add_argument(
        "--port",
        type=int,
        default=5100,
        help="Port on which the web server should run (default: 5100).",
    )
    parser.add_argument(
        "--host",
        type=str,
        default="0.0.0.0",
        help="Host address (default: 0.0.0.0).",
    )
    parser.add_argument(
        "--reload",
        action="store_true",
        help="Optional: enables auto-reload on code changes (development only).",
    )

    args = parser.parse_args()

    uvicorn.run(
        app,
        host=args.host,
        port=args.port,
        reload=args.reload,
        workers=1,
    )