Add metadata and JSON objects

hyperforge/src/hyperforge/api/v1/mcp_content.py

@@ -229,6 +229,19 @@ def convert_arag_answer_to_content(
                 )
             )
 
+        if context.json_objects:
+            for idx, json_obj in enumerate(context.json_objects):
+                contents.append(
+                    EmbeddedResource(
+                        type="resource",
+                        resource=TextResourceContents(
+                            uri=f"rao-response://context/{context.id}/json-objects/{idx}",  # type: ignore[arg-type]
+                            text=json.dumps(pydantic_core.to_jsonable_python(json_obj)),
+                            mimeType="application/json",
+                        ),
+                    )
+                )
+
         if context.image_urls:
             contents.append(
                 EmbeddedResource(
@@ -272,6 +285,7 @@ def convert_arag_answer_to_content(
                 type="text",
                 text="\n".join(parts),
                 annotations=Annotations(audience=["assistant"]),
+                _meta=step.metadata,
             )
         )
 

hyperforge/src/hyperforge/models.py

@@ -122,6 +122,7 @@ class Step(BaseModel):
     input_nuclia_tokens: Optional[float]
     output_nuclia_tokens: Optional[float]
     error: Optional[str] = None
+    metadata: Optional[Dict[str, Any]] = None
 
     def __str__(self):
         return f"({self.timeit:.2f}s) {self.module}: {self.title} \n {self.value} \n {self.reason} \n NT:({self.input_nuclia_tokens}:{self.output_nuclia_tokens})"
@@ -248,6 +249,25 @@ class Answer(BaseModel):
 CONTEXT_PROMPT_TEMPLATE = PROMPT_ENVIRONMENT.from_string(CONTEXT_TEMPLATE)
 
 
+class JSONObject(BaseModel):
+    json_schema: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description="JSON schema that defines the structure of the JSON object.",
+    )
+    json_object: Dict[str, Any] = Field(
+        default_factory=dict,
+        description="The actual JSON object that conforms to the provided JSON schema.",
+    )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default_factory=dict,
+        description="Optional metadata associated with the JSON object.",
+    )
+    id: Optional[str] = Field(
+        default_factory=lambda: uuid.uuid4().hex,
+        description="Unique identifier for this JSON object instance.",
+    )
+
+
 class Context(BaseModel):
     id: str = Field(
         default_factory=lambda: uuid.uuid4().hex,
@@ -260,6 +280,7 @@ class Context(BaseModel):
     images: Dict[str, Image] = Field(default_factory=dict)
     prompts: List[Prompt] = Field(default_factory=list)
     structured: List[str] = Field(default_factory=list)
+    json_objects: List[JSONObject] = Field(default_factory=list)
     source: str
     agent: str
     # XXX: This is not actually a summary, but an answer attempt for now!

hyperforge/tests/models/test_utils.py

@@ -3,7 +3,9 @@
 from enum import Enum
 from typing import List, Optional, Union
 
+import pytest
 from hyperforge.api.utils import to_strict_json_schema
+from hyperforge.models import Context, JSONObject
 from inline_snapshot import snapshot
 from pydantic import BaseModel, Field
 
@@ -347,3 +349,110 @@ def test_nested_inline_ref_expansion() -> None:
             "additionalProperties": False,
         }
     )
+
+
+# ---------------------------------------------------------------------------
+# JSONObject model
+# ---------------------------------------------------------------------------
+
+
+def test_json_object_minimal():
+    """JSONObject can be created with only json_object (all other fields optional)."""
+    obj = JSONObject(json_object={"key": "value"})
+    assert obj.json_object == {"key": "value"}
+    assert obj.json_schema is None
+    assert obj.metadata == {}
+    assert obj.id is not None  # auto-generated uuid
+
+
+def test_json_object_with_schema():
+    """JSONObject stores a JSON schema alongside the object."""
+    schema = {"type": "object", "properties": {"key": {"type": "string"}}}
+    obj = JSONObject(
+        json_schema=schema,
+        json_object={"key": "value"},
+    )
+    assert obj.json_schema == schema
+    assert obj.json_object == {"key": "value"}
+
+
+def test_json_object_with_metadata():
+    """JSONObject stores arbitrary metadata."""
+    obj = JSONObject(
+        json_object={"x": 1},
+        metadata={"source": "test", "confidence": 0.9},
+    )
+    assert obj.metadata == {"source": "test", "confidence": 0.9}
+
+
+def test_json_object_explicit_id():
+    """JSONObject accepts an explicit id and does not override it."""
+    obj = JSONObject(json_object={}, id="my-custom-id")
+    assert obj.id == "my-custom-id"
+
+
+def test_json_object_unique_ids():
+    """Two JSONObjects without explicit ids get different auto-generated ids."""
+    a = JSONObject(json_object={})
+    b = JSONObject(json_object={})
+    assert a.id != b.id
+
+
+def test_json_object_roundtrip():
+    """JSONObject survives model_dump / model_validate round-trip."""
+    obj = JSONObject(
+        json_schema={"type": "object"},
+        json_object={"answer": 42},
+        metadata={"tag": "roundtrip"},
+    )
+    dumped = obj.model_dump()
+    restored = JSONObject.model_validate(dumped)
+    assert restored.json_object == obj.json_object
+    assert restored.json_schema == obj.json_schema
+    assert restored.metadata == obj.metadata
+    assert restored.id == obj.id
+
+
+# ---------------------------------------------------------------------------
+# Context.json_objects field
+# ---------------------------------------------------------------------------
+
+
+def _make_context(**kwargs) -> Context:
+    defaults = dict(
+        original_question_uuid=None,
+        actual_question_uuid=None,
+        question="Q",
+        source="test",
+        agent="test-agent",
+    )
+    return Context(**{**defaults, **kwargs})
+
+
+def test_context_json_objects_defaults_empty():
+    """Context.json_objects is an empty list when not provided."""
+    ctx = _make_context()
+    assert ctx.json_objects == []
+
+
+def test_context_json_objects_stores_items():
+    """Context.json_objects holds JSONObject instances."""
+    objs = [
+        JSONObject(json_object={"a": 1}),
+        JSONObject(json_object={"b": 2}),
+    ]
+    ctx = _make_context(json_objects=objs)
+    assert len(ctx.json_objects) == 2
+    assert ctx.json_objects[0].json_object == {"a": 1}
+    assert ctx.json_objects[1].json_object == {"b": 2}
+
+
+def test_context_json_objects_roundtrip():
+    """Context.json_objects survives model_dump / model_validate round-trip."""
+    objs = [JSONObject(json_object={"key": "val"}, metadata={"m": 1})]
+    ctx = _make_context(json_objects=objs)
+    dumped = ctx.model_dump()
+    restored = Context.model_validate(dumped)
+    assert len(restored.json_objects) == 1
+    assert restored.json_objects[0].json_object == {"key": "val"}
+    assert restored.json_objects[0].metadata == {"m": 1}

hyperforge/tests/test_convert_arag_answer.py

@@ -430,6 +430,76 @@ def test_step_without_optional_fields():
     assert "Planning" in result[0].text
 
 
+def test_step_metadata_defaults_to_none():
+    """Step.metadata is None when not provided (backward compat)."""
+    step = Step(
+        original_question_uuid=None,
+        actual_question_uuid=None,
+        module="smart",
+        title="Planning",
+        timeit=0.5,
+        input_nuclia_tokens=None,
+        output_nuclia_tokens=None,
+        agent_path="agent",
+    )
+    assert step.metadata is None
+
+
+def test_step_metadata_accepts_dict():
+    """Step.metadata stores an arbitrary key/value dict."""
+    meta = {"source": "test-run", "confidence": 0.95, "tags": ["a", "b"]}
+    step = Step(
+        original_question_uuid="q1",
+        actual_question_uuid="q1",
+        module="smart",
+        title="Planning",
+        timeit=1.0,
+        input_nuclia_tokens=5,
+        output_nuclia_tokens=10,
+        agent_path="agent",
+        metadata=meta,
+    )
+    assert step.metadata == meta
+
+
+def test_step_metadata_roundtrip():
+    """Step.metadata survives a model_dump / model_validate round-trip."""
+    meta = {"key": "value", "nested": {"inner": 42}}
+    step = Step(
+        original_question_uuid="q1",
+        actual_question_uuid="q1",
+        module="smart",
+        title="Planning",
+        timeit=1.0,
+        input_nuclia_tokens=5,
+        output_nuclia_tokens=10,
+        agent_path="agent",
+        metadata=meta,
+    )
+    dumped = step.model_dump()
+    assert dumped["metadata"] == meta
+
+    restored = Step.model_validate(dumped)
+    assert restored.metadata == meta
+
+
+def test_step_metadata_none_omitted_in_dump():
+    """When metadata is None the serialised dict contains the key with None value."""
+    step = Step(
+        original_question_uuid=None,
+        actual_question_uuid=None,
+        module="smart",
+        title="Planning",
+        timeit=0.5,
+        input_nuclia_tokens=None,
+        output_nuclia_tokens=None,
+        agent_path="agent",
+    )
+    dumped = step.model_dump()
+    assert "metadata" in dumped
+    assert dumped["metadata"] is None
+
+
 # Exception