fgAI/src/fgai/llm.py

from __future__ import annotations

import json
import os
from urllib import request

from .models import BlockCandidate, Finding


def ollama_summary(
    findings: list[Finding],
    candidates: list[BlockCandidate],
    model: str | None = None,
    *,
    analysis: dict[str, object] | None = None,
    timeout: int | None = None,
) -> str:
    host = os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")
    selected_model = model or os.getenv("OLLAMA_MODEL", "llama3.1")
    selected_timeout = timeout or int(os.getenv("OLLAMA_TIMEOUT", "180"))
    prompt = {
        "analysis": analysis or {},
        "findings": [finding.__dict__ for finding in findings],
        "block_candidates": [
            {"src_ip": candidate.src_ip, "score": candidate.score, "reasons": candidate.reasons}
            for candidate in candidates[:20]
        ],
    }
    body = json.dumps(
        {
            "model": selected_model,
            "stream": False,
            "options": {
                "num_predict": 350,
                "temperature": 0.2,
            },
            "prompt": (
                "You are SignalScope, a local security operations analyst. Analyze only the supplied telemetry. "
                "Never describe the input as JSON, a SIEM object, a dataset, or an array. Never ask the user what to focus on. "
                "Return exactly these short sections: Assessment, Priority entities, Evidence, Recommended next action. "
                "Use actual entity names, stream names, counts, scores, and field deviations from the supplied data. "
                "If evidence is insufficient, say that explicitly and name the missing field or stream. "
                "Do not recommend blocking private/internal client IPs unless the data explicitly proves compromise. "
                f"\n\nTelemetry:\n{json.dumps(prompt)}"
            ),
        }
    ).encode("utf-8")
    req = request.Request(f"{host}/api/generate", data=body, method="POST", headers={"Content-Type": "application/json"})
    with request.urlopen(req, timeout=selected_timeout) as response:
        data = json.loads(response.read().decode("utf-8"))
    return str(data.get("response", "")).strip()


def ollama_dashboard_assessment(analysis: dict[str, object], model: str | None = None, timeout: int | None = None) -> str:
    compact = {
        "summary": analysis.get("summary", {}),
        "anomaly_summary": analysis.get("anomaly_summary", {}),
        "top_anomalies": analysis.get("anomalies", [])[:5],
        "top_recommendations": analysis.get("recommendations", [])[:5],
        "block_candidates": analysis.get("block_candidates", [])[:5],
        "policy_findings": analysis.get("policy_findings", [])[:5],
        "event_context": analysis.get("event_context", {}),
        "diagnostics": analysis.get("diagnostics", {}),
        "capabilities": analysis.get("capabilities", {}),
        "cross_source_correlations": analysis.get("cross_source_correlations", [])[:20],
        "incidents": analysis.get("incidents", [])[:10],
        "profile_suggestions": analysis.get("profile_suggestions", [])[:10],
        "field_deviations": analysis.get("field_deviations", {}),
        "feedback": analysis.get("feedback", []),
    }
    return ollama_summary(
        [],
        [],
        model,
        analysis={
            "task": (
                "Write a concise dashboard analyst note. Compare activity across every listed entity, "
                "identify the most unusual entity or behavior, and state the next investigation step. "
                "Mention policyid=0 as implicit deny/drop, not an editable policy."
            ),
            "data": compact,
        },
        timeout=timeout,
    )


def ollama_profile_advice(suggestions: list[dict[str, object]], model: str | None = None, timeout: int | None = None) -> list[dict[str, object]]:
    host = os.getenv("OLLAMA_HOST", "http://127.0.0.1:11434").rstrip("/")
    selected_model = model or os.getenv("FGAI_PROFILE_ADVISOR_MODEL", "qwen3:8b")
    selected_timeout = timeout or int(os.getenv("FGAI_PROFILE_ADVISOR_TIMEOUT", "120"))
    compact = [
        {
            "stream_id": item.get("stream_id"),
            "stream_name": item.get("stream_name"),
            "events": item.get("events"),
            "common_fields": item.get("common_fields", [])[:20],
            "heuristic_profile": item.get("profile", {}),
        }
        for item in suggestions[:10]
    ]
    body = json.dumps(
        {
            "model": selected_model,
            "stream": False,
            "format": "json",
            "options": {"num_predict": 1200, "temperature": 0.1},
            "prompt": (
                "You are SignalScope's local profile advisor. Infer stream profile mappings from observed field statistics. "
                "Return only valid JSON with this schema: "
                "{\"profiles\":[{\"stream_id\":\"...\",\"entity_fields\":[\"...\"],\"timestamp_field\":\"...\","
                "\"categorical_fields\":[\"...\"],\"numeric_fields\":[\"...\"],\"detectors\":{\"auth_failure\":{\"enabled\":true,\"minimum\":5,\"z_threshold\":3}},"
                "\"reason\":\"short reason\"}]}. "
                "Use only field names present in common_fields or heuristic_profile. Do not include raw message/full_message fields. "
                "Allowed detectors are auth_failure, dns_query, deny_action. Prefer canonical fields such as username, hostname, eventid, srcip, dstip when present. "
                f"\n\nObserved streams:\n{json.dumps(compact, sort_keys=True)}"
            ),
        }
    ).encode("utf-8")
    req = request.Request(f"{host}/api/generate", data=body, method="POST", headers={"Content-Type": "application/json"})
    with request.urlopen(req, timeout=selected_timeout) as response:
        data = json.loads(response.read().decode("utf-8"))
    payload = json.loads(str(data.get("response", "{}")))
    profiles = payload.get("profiles", [])
    return profiles if isinstance(profiles, list) else []