process.py

import os
import json
from typing import Optional
from mistralai import Mistral
from dotenv import load_dotenv
from datetime import datetime
from utils import get_frames, encode_image_bytes

load_dotenv()

api_key = os.environ["MISTRAL_API_KEY"]
model = "ministral-3b-latest"

client = Mistral(api_key=api_key)


def image_inference(image_bytes: bytes):
    messages = [
        {
            "role": "user",
            "content": [
                {
                    "type": "text",
                    "text": "What's in this image? Answer consicely in 2-3 sentences. Be deterministic and do not make assumptions or use uncertain language."
                },
                {
                    "type": "image_url",
                    "image_url": f"data:image/jpeg;base64,{image_bytes}"
                }
            ]
        }
    ]

    chat_response = client.chat.complete(
        model=model,
        messages=messages
    )
    return chat_response.choices[0].message.content


def stream_inference(context: dict, changes: dict,):
    stream_response = client.chat.stream(
        model=model,
        messages=[
            {
                "role": "system",
                "content": f"""You are narrating a video based on detected changes in object state.
Describe ONLY what changed between the last two frames in the object's attributes or state.
DO NOT narrate changes in how the subject is named - those are just linguistic variations, not real changes.
Be concise and deterministic. Do not include formatting. Keep the response to maximum of 2 sentences. Use plain paragraphs.
Do not repeat information already provided in the context.
Focus on the actual change, not restating unchanged facts. 
Mention concrete changes only, do not mention subtle phrasing changes""",
            },
            {
                "role": "user",
                "content": f"{json.dumps(context)}\nDetected changes: {json.dumps(changes)}"
            }
        ]
    )
    accum = ""
    accum_size = 10
    t_count = 0
    for chunk in stream_response:
        content = chunk.data.choices[0].delta.content
        if content:
            accum += content
            t_count += 1

            if t_count >= accum_size:
                print(accum, end="", flush=True)
                accum = ""
                t_count = 0
    if accum:
        print(accum, end="", flush=True)


def get_structured_state(image_desc: str):
    tools = [
        {
            "type": "function",
            "function": {
                "name": "get_state",
                "description": "Extracts structured state from a descriptive sentence",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "subject": {
                            "type": "string",
                            "description": "Primary object or entity in the image"
                        },
                        "attributes": {
                            "type": "object",
                            "properties": {
                                "color": {
                                    "type": "string"
                                },
                                "shape": {
                                    "type": "string"
                                },
                                "material": {
                                    "type": "string"
                                },
                                "size": {
                                    "type": "string",
                                    "description": "Relative size perception (small, medium, large, growing, shrinking)"
                                },
                                "count": {
                                    "type": "integer",
                                    "description": "Number of primary subjects visible in frame"
                                }
                            }
                        },
                        "state": {
                            "type": "object",
                            "properties": {
                                "motion": {
                                    "type": "object",
                                    "properties": {
                                        "status": {
                                            "type": "string",
                                            "enum": [
                                                "stationary", "moving", "rotating", "unknown"
                                            ]
                                        },
                                        "details": {
                                            "type": "string",
                                            "description": "Details of what the subject is doing and their position"
                                        }
                                    }
                                },
                                "orientation": {
                                    "type": "string",
                                    "description": "Orientation or facing direction (upright, tilted, sideways, etc.)"
                                },
                                "position": {
                                    "type": "string",
                                    "description": "Position in frame (left, center, right, top, bottom, etc.)"
                                },
                                "interaction": {
                                    "type": "string",
                                    "description": "Interaction state with other objects (isolated, touching, holding, connected, etc.)"
                                }
                            }
                        }
                    }
                },
                "strict": True
            }
        }
    ]
    messages = [
        {"role": "system", "content": """Use the get_state tool to extract structured information from a given sentence. 
    Extract only directly observable properties.
    Do not infer intent, cause, or future action.
    Use coarse, stable labels.
    Use short canonical labels.
    Do not include adjectives, qualifiers, or brand-style phrasing.
    If unsure, output unknown.
    """},
        {"role": "user", "content": image_desc}
    ]
    response = client.chat.complete(
        messages=messages,
        temperature=0,
        model=model,
        tools=tools)
    return response.choices[0].message.tool_calls[0].function.arguments


def diff_states(prev_state, curr_state):
    """
    Diff two states and return changes in event-relevant fields.

    Supports both 2-level paths like ("attributes", "color")
    and nested 3-level paths like ("state", "motion", "status").
    """
    EVENT_FIELDS = [
        ("attributes", "color"),
        # ("attributes", "size"),
        ("attributes", "count"),
        ("state", "motion", "status"),  # Nested path for motion status
        ("state", "position"),
        ("state", "orientation"),
        ("state", "interaction"),
    ]

    if prev_state is None:
        return None

    changes = {}

    for field_path in EVENT_FIELDS:
        # Navigate through the path
        prev_val = prev_state
        curr_val = curr_state

        for key in field_path:
            prev_val = prev_val.get(key, {}) if isinstance(prev_val, dict) else None
            curr_val = curr_val.get(key, {}) if isinstance(curr_val, dict) else None

        # Ignore missing values
        if prev_val is None or curr_val is None:
            continue

        # Ignore if either value is unknown (not a real change)
        if prev_val == "unknown" or curr_val == "unknown":
            continue

        # Record actual changes
        if prev_val != curr_val:
            field_name = ".".join(field_path)
            changes[field_name] = {
                "from": prev_val,
                "to": curr_val
            }

    return changes


def get_stream(video_path: str, frame_step: int = 10, limit: Optional[int] = None):
    """
    Process video frames and emit streaming narration of state changes.

    Args:
        video_path: Path to the video file
        frame_step: Process every Nth frame
        limit: Maximum number of frames to process (None = all)

    Returns:
        VideoContext: Contains canonical_context, state_history, and change_history
    """
    frames = get_frames(video_path, frame_step)
    frames_to_process = frames if limit is None else frames[:limit]
    prev_state = None
    canonical_context = None
    state_history = []
    change_history = []

    for frame_dict in frames_to_process:
        image_bytes = encode_image_bytes(frame_dict["image"])
        current_response = image_inference(image_bytes)
        curr_state = json.loads(get_structured_state(current_response))

        # Initialize canonical context on first frame
        if canonical_context is None:
            canonical_subject = curr_state["subject"]
            canonical_context = {
                "subject": canonical_subject,
                "attributes": {},
                "state": {}
            }
            color = curr_state.get("attributes", {}).get("color")
            motion = curr_state.get("state", {}).get("motion")

            if color and color != "unknown":
                canonical_context["attributes"]["color"] = color
            if motion and motion != "unknown":
                canonical_context["state"]["motion"] = motion

            state_history.append(curr_state)
            prev_state = curr_state
            continue

        # Diff states
        changes = diff_states(prev_state, curr_state)

        if changes:
            for path, change in changes.items():
                # Handle variable-length paths (e.g., "state.motion.status")
                parts = path.split(".")

                # Navigate to the correct nested location
                target = canonical_context
                for part in parts[:-1]:
                    if part not in target:
                        target[part] = {}
                    target = target[part]

                # Set the final value
                target[parts[-1]] = change["to"]

            change_history.append(changes)
            stream_inference(canonical_context, changes)
            print()  # Newline after each narration

        state_history.append(curr_state)
        prev_state = curr_state

    return {
        "canonical_context": canonical_context,
        "state_history": state_history,
        "change_history": change_history
    }


def answer_question(question: str, video_context: dict) -> str:
    """
    Answer questions about the processed video using accumulated context.

    Args:
        question: User's question about the video
        video_context: Dict containing canonical_context, state_history, change_history

    Returns:
        str: The answer to the question
    """
    context_summary = f"""Video Context:
Subject: {video_context['canonical_context']['subject']}
Final State: {json.dumps(video_context['canonical_context'], indent=2)}

State History ({len(video_context['state_history'])} frames):
{json.dumps(video_context['state_history'], indent=2)}

Detected Changes ({len(video_context['change_history'])} events):
{json.dumps(video_context['change_history'], indent=2)}
"""

    response = client.chat.complete(
        model=model,
        messages=[
            {
                "role": "system",
                "content": """You are a video understanding assistant. You have access to:
1. The complete state history of objects in a video
2. All detected state changes
3. The final canonical state

Answer questions precisely based on this data. Be concise and factual.
If the information isn't in the context, say so.

Answer in a single text paragraph with no formatting.
Do not make assumptions.
"""
            },
            {
                "role": "user",
                "content": f"{context_summary}\n\nQuestion: {question}"
            }
        ]
    )
    return response.choices[0].message.content


def interactive_qa(video_context: dict):
    """
    Interactive Q&A loop for querying processed video context.

    Args:
        video_context: Dict containing canonical_context, state_history, change_history
    """
    print("\n" + "="*60)
    print("Interactive Q&A Mode")
    print("="*60)
    print("Ask questions about the video. Type 'exit' or 'quit' to end.\n")

    while True:
        try:
            question = input("\nQuestion: ").strip()

            if question.lower() in ["exit", "quit", "q"]:
                print("Exiting Q&A mode.")
                break

            if not question:
                continue

            answer = answer_question(question, video_context)
            print(f"\nAnswer: {answer}")

        except KeyboardInterrupt:
            print("\n\nExiting Q&A mode.")
            break
        except Exception as e:
            print(f"\nError: {e}")


if __name__ == "__main__":
    # Phase 1: Process video and stream narrations
    print("Processing video...\n")
    video_context = get_stream("./input/sample2.mp4", frame_step=10)

    # Phase 2: Interactive Q&A
    if video_context and video_context.get("canonical_context"):
        interactive_qa(video_context)
    else:
        print("\nNo video context available for Q&A.")