core.py

"""
Shared logic core for the VLPS demo — model/index resources and the business
functions behind every feature, returning plain Python data (no UI framework).

Both the FastAPI site (server.py) and the Gradio app (app.py) import from here,
so there is a single source of truth.
"""
import io
import os
import sys
from collections import OrderedDict
from functools import lru_cache

os.environ.setdefault("KMP_DUPLICATE_LIB_OK", "TRUE")
sys.path.append(os.path.abspath("src"))

from dotenv import load_dotenv
from PIL import Image

from vlps.opensearch_config import get_client
from vlps.data import ask_lvlm, load_siglip_model, embed_text_siglip
from vlps import ocr, textvqa

load_dotenv(".env", override=True)

BASE_URL = os.getenv("VLLM_BASE_URL", "https://api.novasearch.org/gemma4/v1")
API_KEY  = os.getenv("VLLM_API_KEY", "nova-vl")
MODEL    = os.getenv("VLLM_MODEL", "google/gemma-4-31B-it")
USER     = os.getenv("OPENSEARCH_USER", "uservl10")

INDEX_BGE   = f"{USER}_coco_bge"
INDEX_MULTI = f"{USER}_coco_multimodal"
INDEX_OCR   = f"{USER}_textvqa"


# ---------------------------------------------------------------------------
# Lazy, cached resources
# ---------------------------------------------------------------------------

@lru_cache(maxsize=1)
def os_client():
    return get_client()


@lru_cache(maxsize=1)
def llm():
    from openai import OpenAI
    return OpenAI(base_url=BASE_URL, api_key=API_KEY)


@lru_cache(maxsize=1)
def bge():
    from sentence_transformers import SentenceTransformer
    return SentenceTransformer("BAAI/bge-base-en-v1.5")


@lru_cache(maxsize=1)
def siglip():
    return load_siglip_model()  # (processor, model, device)


@lru_cache(maxsize=1)
def ocr_reader_en():
    return ocr.get_reader(("en",))


@lru_cache(maxsize=1)
def ocr_reader_pt():
    return ocr.get_reader(("pt", "en"))


def coco_image_url(image_id) -> str:
    """Public COCO image URL.

    HuggingFaceM4/COCO uses 2014 cocoids, so images live in val2014 with the
    split-prefixed filename. The S3 path-style host has a valid TLS cert (the bare
    images.cocodataset.org domain does not).
    """
    return f"https://s3.amazonaws.com/images.cocodataset.org/val2014/COCO_val2014_{int(image_id):012d}.jpg"


@lru_cache(maxsize=512)
def fetch_image(url: str):
    import requests
    try:
        r = requests.get(url, timeout=20)
        r.raise_for_status()
        return Image.open(io.BytesIO(r.content)).convert("RGB")
    except Exception:  # noqa: BLE001
        return None


@lru_cache(maxsize=1)
def pt_data():
    """(records, {image_id: PIL}) for the student Portuguese set, or ([], {})."""
    try:
        return textvqa.load_pt_dataset("data/pt_textvqa")
    except FileNotFoundError:
        return [], {}


# ---------------------------------------------------------------------------
# Feature functions (return plain data)
# ---------------------------------------------------------------------------

def search_images(query: str, mode: str = "knn", top_k: int = 6) -> list:
    """Phase 1 retrieval. mode in {'bm25','knn','crossmodal'}. Returns list of dicts."""
    client = os_client()
    if mode == "bm25":
        body = {"size": top_k, "query": {"match": {"caption": query}}}
        index = INDEX_BGE
    elif mode == "crossmodal":
        proc, model, dev = siglip()
        vec = embed_text_siglip(query, proc, model, dev)
        body = {"size": top_k, "query": {"knn": {"image_vec": {"vector": vec, "k": top_k}}}}
        index = INDEX_MULTI
    else:  # knn (BGE)
        vec = bge().encode(query, normalize_embeddings=True).tolist()
        body = {"size": top_k, "query": {"knn": {"caption_vec": {"vector": vec, "k": top_k}}}}
        index = INDEX_BGE

    resp = client.search(index=index, body=body)
    out = []
    for h in resp["hits"]["hits"]:
        s = h["_source"]
        out.append({
            "image_id": s["image_id"],
            "caption": s.get("caption", ""),
            "score": round(h["_score"], 4),
            "image_url": coco_image_url(s["image_id"]),
        })
    return out


def visual_qa(image: Image.Image, question: str, system_prompt: str = "") -> str:
    return ask_lvlm(llm(), image, question, model=MODEL, system_prompt=system_prompt or None)


def ocr_qa(image: Image.Image, question: str, lang: str = "en") -> dict:
    reader = ocr_reader_pt() if lang == "pt" else ocr_reader_en()
    res = ocr.ocr_image(image, reader=reader)
    return {
        "ocr_text": res["text"],
        "answer_ocr": textvqa.answer_with_ocr(llm(), res["text"], question, model=MODEL),
        "answer_vision": textvqa.answer_with_vision(llm(), image, question, model=MODEL),
    }


def _docs_to_items(sources: list, dataset: str) -> list:
    items = []
    for s in sources:
        iid = s["image_id"]
        url = s.get("image_url") or ""
        if dataset == "pt" and not url:
            url = f"/pt-image/{iid}"
        items.append({
            "image_id": iid,
            "image_url": url,
            "question": s.get("question", ""),
            "answers": s.get("answers", []),
            "ocr_text": s.get("ocr_text", ""),
            "language": s.get("language", ""),
            "dataset": s.get("dataset", dataset),
        })
    return items


def dataset_browse(dataset: str = "textvqa", limit: int = 12) -> list:
    resp = os_client().search(
        index=INDEX_OCR,
        body={"size": int(limit), "query": {"term": {"dataset": dataset}}},
    )
    return _docs_to_items([h["_source"] for h in resp["hits"]["hits"]], dataset)


def dataset_search(dataset: str, query: str, k: int = 6) -> list:
    hits = textvqa.retrieve_by_ocr_text(os_client(), query, INDEX_OCR, top_k=int(k), dataset=dataset)
    return _docs_to_items(hits, dataset)


# What the agent does during a run (so the UI can show it). Reset per request.
_agent_seen: list = []
_agent_steps: list = []


def _note_image(image_id, caption=""):
    _agent_seen.append({"image_id": str(image_id), "caption": caption,
                        "image_url": coco_image_url(image_id)})


def _note_step(tool: str, **args):
    _agent_steps.append({"tool": tool, "args": args})


@lru_cache(maxsize=1)
def _agent_model():
    from smolagents import OpenAIServerModel
    return OpenAIServerModel(model_id=MODEL, api_base=BASE_URL, api_key=API_KEY)


@lru_cache(maxsize=1)
def _agent_tools():
    from smolagents import tool
    from vlps.data import retrieve_by_text_knn

    @tool
    def retrieve_images_by_text(query: str, top_k: int = 3) -> list:
        """
        Find COCO images matching a natural-language description (BGE semantic search).
        Returns a list of {image_id, caption} dicts.

        Args:
            query: Natural-language description of the images to find.
            top_k: Number of images to return.
        """
        _note_step("retrieve_images_by_text", query=query, top_k=top_k)
        hits = retrieve_by_text_knn(os_client(), query, INDEX_BGE, bge(), top_k)
        out = [{"image_id": h["image_id"], "caption": h.get("caption", "")} for h in hits]
        for h in out:
            _note_image(h["image_id"], h["caption"])
        return out

    @tool
    def answer_about_image(image_id: str, question: str) -> str:
        """
        Answer a visual question about a specific COCO image with the vision model.
        Call retrieve_images_by_text first to obtain a relevant image_id.

        Args:
            image_id: COCO image id to inspect.
            question: The visual question to answer.
        """
        _note_step("answer_about_image", image_id=image_id)
        img = fetch_image(coco_image_url(image_id))
        if img is None:
            return f"Image {image_id} not found."
        _note_image(image_id)
        return ask_lvlm(llm(), img, question, model=MODEL,
                        system_prompt="You are an expert visual analyst. Answer concisely.")

    return [retrieve_images_by_text, answer_about_image]


# One agent per browser session, so conversations stay isolated and remember context.
_AGENTS: "OrderedDict[str, object]" = OrderedDict()
_MAX_SESSIONS = 64


def _get_agent(session_id: str):
    """Return (agent, is_new) for a session, evicting the oldest beyond the cap."""
    from smolagents import ToolCallingAgent
    if session_id in _AGENTS:
        _AGENTS.move_to_end(session_id)
        return _AGENTS[session_id], False
    ag = ToolCallingAgent(tools=_agent_tools(), model=_agent_model(), max_steps=6)
    _AGENTS[session_id] = ag
    if len(_AGENTS) > _MAX_SESSIONS:
        _AGENTS.popitem(last=False)
    return ag, True


def reset_agent(session_id: str):
    """Drop a session's conversation memory."""
    _AGENTS.pop(session_id, None)


def agent_answer(message: str, session_id: str = "") -> dict:
    """
    Run the agent and return its answer plus any images it looked at (deduped).
    With a session_id the conversation keeps context across turns; the first turn
    starts fresh, later turns retain memory.
    """
    _agent_seen.clear()
    _agent_steps.clear()
    if not session_id:
        session_id = "anon"
    ag, is_new = _get_agent(session_id)
    answer = str(ag.run(message, reset=is_new))  # first turn fresh, later turns keep memory
    seen, images = set(), []
    for it in _agent_seen:
        if it["image_id"] not in seen:
            seen.add(it["image_id"])
            images.append(it)
    return {"answer": answer, "images": images, "tools": list(_agent_steps)}