refactor(core): consolidate backend raw responses onto mot.raw namespace

cli/serve/utils.py

@@ -35,31 +35,30 @@ def extract_finish_reason(output: Any) -> FinishReason:
         "function_call",
     }
 
-    # Try to get finish_reason from the response metadata
-    # Different backends store this in different places
-    if hasattr(output, "_meta") and output._meta:
-        # Ollama backend stores response in chat_response with done_reason field
-        # (ollama.ChatResponse object with done_reason attribute)
-        chat_response = output._meta.get("chat_response")
-        if chat_response and hasattr(chat_response, "done_reason"):
-            done_reason = chat_response.done_reason
+    # Try to get finish_reason from the backend-native response on mot.raw.
+    # Different backends store this in different places; switch on mot.raw.provider.
+    raw = getattr(output, "raw", None)
+    if raw is not None:
+        provider = raw.provider
+        response = raw.response
+
+        if provider == "ollama" and response is not None:
+            # ollama.ChatResponse object with done_reason attribute.
+            done_reason = getattr(response, "done_reason", None)
             if done_reason in valid_reasons:
                 return done_reason
 
-        # OpenAI backend stores full response dict in oai_chat_response
-        # (from chunk.model_dump() which includes choices array)
-        oai_response = output._meta.get("oai_chat_response")
-        if oai_response and isinstance(oai_response, dict):
-            choices = oai_response.get("choices", [])
+        elif provider in ("openai", "watsonx", "litellm") and isinstance(
+            response, dict
+        ):
+            # Chat path: full response dict, finish_reason nested under choices[0].
+            choices = response.get("choices", [])
             if choices and len(choices) > 0:
                 finish_reason = choices[0].get("finish_reason")
                 if finish_reason in valid_reasons:
                     return finish_reason
-
-        # LiteLLM backend stores response dict in litellm_chat_response
-        litellm_response = output._meta.get("litellm_chat_response")
-        if litellm_response and isinstance(litellm_response, dict):
-            finish_reason = litellm_response.get("finish_reason")
+            # Raw-completion path: single choice dict, finish_reason at top level.
+            finish_reason = response.get("finish_reason")
             if finish_reason in valid_reasons:
                 return finish_reason
 

mellea/backends/huggingface.py

@@ -1233,7 +1233,7 @@ async def processing(
             mot._underlying_value += chunk
         elif isinstance(chunk, GenerateDecoderOnlyOutput):
             # Otherwise, it's a non-streaming request. Decode it here.
-            mot._meta["hf_output"] = chunk
+            mot.raw.response = chunk
             mot._underlying_value += cast(
                 str,
                 self._tokenizer.decode(
@@ -1269,18 +1269,18 @@ class used during generation, if any.
             input_ids: The prompt token IDs; used to compute token counts and for
                 KV cache bookkeeping.
         """
-        if mot._meta.get("hf_output", None) is None:
+        if mot.raw.response is None:
             if mot._generate_extra is not None:
                 full_output = await mot._generate_extra
                 assert isinstance(full_output, GenerateDecoderOnlyOutput)
-                mot._meta["hf_output"] = full_output
+                mot.raw.response = full_output
 
         # The ModelOutputThunk must be computed by this point.
         assert mot.value is not None
 
-        # Store KV cache in LRU separately (not in mot._meta) to enable proper cleanup on eviction.
+        # Store KV cache in LRU separately (not on the MOT) to enable proper cleanup on eviction.
         # This prevents GPU memory from being held by ModelOutputThunk references.
-        hf_output = mot._meta.get("hf_output", None)
+        hf_output = mot.raw.response
         if (
             self._use_caches
             and isinstance(hf_output, GenerateDecoderOnlyOutput)
@@ -1320,7 +1320,7 @@ class used during generation, if any.
         )
 
         # Derive token counts from the output sequences (HF models have no usage object).
-        hf_output = mot._meta.get("hf_output")
+        hf_output = mot.raw.response
         n_prompt, n_completion = None, None
         if isinstance(hf_output, GenerateDecoderOnlyOutput):
             try:
@@ -1375,20 +1375,21 @@ class used during generation, if any.
         # Populate model and provider metadata
         mot.generation.model = self._model_id
         mot.generation.provider = self._provider
+        mot.raw.provider = self._provider
 
-        # When caching is disabled, clear hf_output from meta to free GPU memory.
+        # When caching is disabled, clear hf_output from raw to free GPU memory.
         # The sequences tensor is on GPU and accumulates if not cleared.
         if not self._use_caches and isinstance(
-            mot._meta.get("hf_output"), GenerateDecoderOnlyOutput
+            mot.raw.response, GenerateDecoderOnlyOutput
         ):
             import gc
 
-            hf_out = mot._meta["hf_output"]
+            hf_out = mot.raw.response
             if hasattr(hf_out, "sequences") and hf_out.sequences is not None:
                 del hf_out.sequences
             if hasattr(hf_out, "scores") and hf_out.scores is not None:
                 del hf_out.scores
-            del mot._meta["hf_output"]
+            mot.raw.response = None
 
             # Force Python GC and return CUDA memory to device
             gc.collect()
@@ -1521,6 +1522,7 @@ async def _generate_from_raw(
             result.generation.usage = per_mot_usage
             result.generation.model = self._model_id
             result.generation.provider = self._provider
+            result.raw.provider = self._provider
 
             action = actions[i]
             result.parsed_repr = (

mellea/backends/litellm.py

@@ -29,6 +29,7 @@
     MelleaLogger,
     ModelOutputThunk,
     ModelToolCall,
+    RawProviderResponse,
 )
 from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
@@ -503,10 +504,8 @@ async def processing(
             if content_chunk is not None:
                 mot._underlying_value += content_chunk
 
-            # Store the full response (includes usage) as a dict
-            mot._meta["litellm_full_response"] = chunk.model_dump()
-            # Also store just the choice for backward compatibility
-            mot._meta["litellm_chat_response"] = chunk.choices[0].model_dump()
+            # Store the full response (includes usage) as a dict.
+            mot.raw.response = chunk.model_dump()
 
         elif isinstance(chunk, litellm.ModelResponseStream):  # type: ignore
             message_delta = chunk.choices[0].delta
@@ -522,15 +521,13 @@ async def processing(
             if content_chunk is not None:
                 mot._underlying_value += content_chunk
 
-            if mot._meta.get("litellm_chat_response_streamed", None) is None:
-                mot._meta["litellm_chat_response_streamed"] = []
-            mot._meta["litellm_chat_response_streamed"].append(
-                chunk.choices[0].model_dump()
-            )
+            if mot.raw.streamed_chunks is None:
+                mot.raw.streamed_chunks = []
+            mot.raw.streamed_chunks.append(chunk.choices[0].model_dump())
 
-            # Store usage information from the chunk if available (typically in the last chunk)
+            # Usage arrives on its own chunk (typically the last); record it now.
             if hasattr(chunk, "usage") and chunk.usage is not None:
-                mot._meta["litellm_streaming_usage"] = chunk.usage.model_dump()
+                mot.generation.usage = chunk.usage.model_dump()
 
     async def post_processing(
         self,
@@ -555,16 +552,13 @@ async def post_processing(
                 `None` if reasoning mode was not enabled.
             _format: The structured output format class used during generation, if any.
         """
-        # Reconstruct the chat_response from chunks if streamed.
-        streamed_chunks = mot._meta.get("litellm_chat_response_streamed", None)
-        if streamed_chunks is not None:
+        # Reconstruct the top-level response from chunks if streamed.
+        if mot.raw.streamed_chunks is not None:
             # Must handle ollama differently due to: https://github.com/BerriAI/litellm/issues/14579.
             # Check that we are targeting ollama with the model_id prefix litellm uses.
-            separate_tools = False
-            if "ollama" in self._model_id.split("/")[0]:
-                separate_tools = True
-            mot._meta["litellm_chat_response"] = chat_completion_delta_merge(
-                streamed_chunks, force_all_tool_calls_separate=separate_tools
+            separate_tools = "ollama" in self._model_id.split("/")[0]
+            mot.raw.response = chat_completion_delta_merge(
+                mot.raw.streamed_chunks, force_all_tool_calls_separate=separate_tools
             )
 
         assert mot._action is not None, (
@@ -577,9 +571,16 @@ async def post_processing(
         # OpenAI-like streamed responses potentially give you chunks of tool calls.
         # As a result, we have to store data between calls and only then
         # check for complete tool calls in the post_processing step.
-        tool_chunk = extract_model_tool_requests(
-            tools, mot._meta["litellm_chat_response"]
+        # Non-streaming stores a top-level response (index into choices); streaming
+        # stores the already-merged choice dict (use directly).
+        response = mot.raw.response
+        assert response is not None
+        choice_response = (
+            response["choices"][0]
+            if isinstance(response, dict) and "choices" in response
+            else response
         )
+        tool_chunk = extract_model_tool_requests(tools, choice_response)
         if tool_chunk is not None:
             if mot.tool_calls is None:
                 mot.tool_calls = {}
@@ -593,7 +594,7 @@ async def post_processing(
         generate_log.backend = f"litellm::{self.model_id!s}"
         generate_log.model_options = mot._model_options
         generate_log.date = datetime.datetime.now()
-        generate_log.model_output = mot._meta["litellm_chat_response"]
+        generate_log.model_output = response
         generate_log.extra = {
             "format": _format,
             "tools_available": tools,
@@ -604,25 +605,18 @@ async def post_processing(
         generate_log.result = mot
         mot._generate_log = generate_log
 
-        # Extract token usage from full response dict or streaming usage
-        full_response = mot._meta.get("litellm_full_response")
-        usage = full_response.get("usage") if isinstance(full_response, dict) else None
-
-        # For streaming responses, usage is stored separately
-        if usage is None:
-            usage = mot._meta.get("litellm_streaming_usage")
-
-        # Populate standardized usage field (LiteLLM uses OpenAI format)
-        if usage:
+        # Non-streaming carries usage on the response; streaming already set it.
+        if usage := response.get("usage"):
             mot.generation.usage = usage
 
         # Populate model and provider metadata
         mot.generation.model = self._model_id
         mot.generation.provider = self._provider
+        mot.raw.provider = self._provider
 
         # Populate response-side metadata for telemetry
-        if isinstance(full_response, dict):
-            populate_response_metadata_openai_shape(mot, full_response)
+        if isinstance(response, dict):
+            populate_response_metadata_openai_shape(mot, response)
 
     @staticmethod
     def _extract_tools(
@@ -730,7 +724,9 @@ async def _generate_from_raw(
             output._context = None  # There is no context for generate_from_raw for now
             output._action = action
             output._model_options = model_opts
-            output._meta = {"litellm_chat_response": res.model_dump()}
+            output.raw = RawProviderResponse(
+                provider=self._provider, response=res.model_dump()
+            )
             output.generation.model = self._model_id
             output.generation.provider = self._provider
 

mellea/backends/ollama.py

@@ -23,6 +23,7 @@
     MelleaLogger,
     ModelOutputThunk,
     ModelToolCall,
+    RawProviderResponse,
 )
 from ..core.base import AbstractMelleaTool
 from ..formatters import ChatFormatter, TemplateFormatter
@@ -608,9 +609,9 @@ async def _generate_from_raw(
                         "completion_tokens": n_out,
                         "total_tokens": n_in + n_out,
                     }
-                result = ModelOutputThunk(
-                    value=response.response,
-                    meta={"generate_response": response.model_dump()},
+                result = ModelOutputThunk(value=response.response)
+                result.raw = RawProviderResponse(
+                    provider=self._provider, response=response.model_dump()
                 )
             result.generation.usage = per_mot_usage
             result.generation.model = self._model_id
@@ -684,9 +685,9 @@ async def processing(
     ):
         """Accumulate text and tool calls from a single Ollama ChatResponse chunk.
 
-        Called for each streaming or non-streaming ``ollama.ChatResponse``. Also
+        Called for each streaming or non-streaming `ollama.ChatResponse`. Also
         extracts tool call requests inline and merges the chunk into the running
-        aggregated response stored in ``mot._meta["chat_response"]``.
+        aggregated response stored in `mot.raw.response`.
 
         Args:
             mot (ModelOutputThunk): The output thunk being populated.
@@ -751,7 +752,7 @@ async def post_processing(
         generate_log.backend = f"ollama::{self._model_id}"
         generate_log.model_options = mot._model_options
         generate_log.date = datetime.datetime.now()
-        generate_log.model_output = mot._meta["chat_response"]
+        generate_log.model_output = mot.raw.response
         generate_log.extra = {
             "format": _format,
             "thinking": mot._model_options.get(ModelOption.THINKING, None),
@@ -766,7 +767,7 @@ async def post_processing(
         mot._generate = None
 
         # Extract token counts from response
-        response = mot._meta.get("chat_response")
+        response = mot.raw.response
         prompt_tokens = (
             getattr(response, "prompt_eval_count", None) if response else None
         )
@@ -783,6 +784,7 @@ async def post_processing(
         # Populate model and provider metadata
         mot.generation.model = self._model_id
         mot.generation.provider = self._provider
+        mot.raw.provider = self._provider
 
         # Populate response-side metadata for telemetry
         if response is not None:
@@ -798,11 +800,11 @@ def chat_response_delta_merge(mot: ModelOutputThunk, delta: ollama.ChatResponse)
         mot: the ModelOutputThunk that the deltas are being used to populated.
         delta: the most recent ollama ChatResponse.
     """
-    if mot._meta.get("chat_response", None) is None:
-        mot._meta["chat_response"] = delta
+    if mot.raw.response is None:
+        mot.raw.response = delta
         return  # Return early, no need to merge.
 
-    merged: ollama.ChatResponse = mot._meta["chat_response"]
+    merged: ollama.ChatResponse = mot.raw.response
     if not merged.done:
         merged.done = delta.done
     if merged.done_reason is None: