feat: making generate_from_raw public (#219)

* feat: making generate_from_raw public in openai

* feat: making generate_from_raw public in ollama

* feat: making generate_from_raw public in ollama

* feat: making generate_from_raw public in hf, watsonx

* feat: making generate_from_raw public in vllm

* adding warning to hf backend

* attempt at standardizing usage metrics

Author: avinash2692

mellea/backends/__init__.py

@@ -58,19 +58,21 @@ def generate_from_context(
         ...
 
     @abc.abstractmethod
-    def _generate_from_raw(
+    def generate_from_raw(
         self,
         actions: list[Component | CBlock],
+        ctx: Context,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
-        generate_logs: list[GenerateLog] | None = None,
+        tool_calls: bool = False,
     ) -> list[ModelOutputThunk]:
         """Generates a model output from the provided input. Does not use context or templates.
 
         Args:
             actions: list of actions to generate responses for. Each action is separate.
+            ctx: context passed to generation. Currently not used in generate_from_raw
             format: A response format to used for structured outputs / constrained decoding. Note: some backends do not support this parameter. They will log warnings and continue to generate.
             model_options: Any model options to upsert into the defaults for this call.
-            generate_logs: a `GenerateLog` instance to add log information to.
+            tool_calls: Always set to false unless supported by backend.
         """

mellea/backends/huggingface.py

@@ -497,15 +497,21 @@ async def post_processing(
 
         mot._generate_log = generate_log
 
-    def _generate_from_raw(
+    def generate_from_raw(
         self,
         actions: list[Component | CBlock],
+        ctx: Context,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
-        generate_logs: list[GenerateLog] | None = None,
+        tool_calls: bool = False,
     ) -> list[ModelOutputThunk]:
         """Generate using the completions api. Gives the input provided to the model without templating."""
+        if tool_calls:
+            FancyLogger.get_logger().warning(
+                "The raw endpoint does not support tool calling at the moment."
+            )
+
         model_opts = self._simplify_and_merge(model_options)
         seed = model_opts.get(ModelOption.SEED, None)
         if seed is not None:
@@ -561,28 +567,34 @@ def _generate_from_raw(
             sequences_to_decode, skip_special_tokens=True
         )
 
-        results = [
-            ModelOutputThunk(value=decoded_result) for decoded_result in decoded_results
-        ]
+        results = []
+        for i, decoded_result in enumerate(decoded_results):
+            n_prompt_tokens = inputs["input_ids"][i].size(0)  # type: ignore
+            n_completion_tokens = len(sequences_to_decode[i])
+            result = ModelOutputThunk(
+                value=decoded_result,
+                meta={
+                    "usage": {
+                        "prompt_tokens": n_prompt_tokens,  # type: ignore
+                        "completion_tokens": n_completion_tokens,
+                        "total_tokens": n_prompt_tokens + n_completion_tokens,
+                    }
+                },
+            )
 
-        for i, result in enumerate(results):
             self.formatter.parse(actions[i], result)
 
-        if generate_logs is not None:
-            assert isinstance(generate_logs, list)
-            date = datetime.datetime.now()
-
-            for i in range(len(prompts)):
-                generate_log = GenerateLog()
-                generate_log.prompt = prompts[i]
-                generate_log.backend = f"hf::{self.model_id!s}"
-                generate_log.model_options = model_opts
-                generate_log.date = date
-                generate_log.model_output = decoded_results
-                generate_log.extra = {"format": format, "seed": seed}
-                generate_log.action = actions[i]
-                generate_log.result = results[i]
-                generate_logs.append(generate_log)
+            generate_log = GenerateLog()
+            generate_log.prompt = self.formatter.print(actions[i])
+            generate_log.backend = f"hf::{self.model_id!s}"
+            generate_log.model_options = model_opts
+            generate_log.date = datetime.datetime.now()
+            generate_log.model_output = decoded_result
+            generate_log.extra = {"format": format, "seed": seed}
+            generate_log.action = actions[i]
+
+            result._generate_log = generate_log
+            results.append(result)
 
         return results
 

mellea/backends/litellm.py

@@ -475,13 +475,14 @@ def _extract_tools(
             FancyLogger.get_logger().info(f"Tools for call: {tools.keys()}")
         return tools
 
-    def _generate_from_raw(
+    def generate_from_raw(
         self,
         actions: list[Component | CBlock],
+        ctx: Context,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
-        generate_logs: list[GenerateLog] | None = None,
+        tool_calls: bool = False,
     ) -> list[ModelOutputThunk]:
         """Generate using the completions api. Gives the input provided to the model without templating."""
         raise NotImplementedError("This method is not implemented yet.")

mellea/backends/ollama.py

@@ -385,19 +385,24 @@ def generate_from_chat_context(
 
         return output
 
-    def _generate_from_raw(
+    def generate_from_raw(
         self,
         actions: list[Component | CBlock],
+        ctx: Context,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
-        generate_logs: list[GenerateLog] | None = None,
+        tool_calls: bool = False,
     ) -> list[ModelOutputThunk]:
         """Generate using the generate api. Gives the input provided to the model without templating."""
         if len(actions) > 1:
             FancyLogger.get_logger().info(
                 "Ollama doesn't support batching; will attempt to process concurrently."
             )
+        if tool_calls:
+            FancyLogger.get_logger().warning(
+                "The completion endpoint does not support tool calling at the moment."
+            )
 
         model_opts = self._simplify_and_merge(model_options)
 
@@ -438,32 +443,41 @@ async def get_response():
             else:
                 result = ModelOutputThunk(
                     value=response.response,
-                    meta={"generate_response": response.model_dump()},
+                    meta={
+                        "generate_response": response.model_dump(),
+                        "usage": {
+                            "completion_tokens": response.eval_count,
+                            "prompt_tokens": response.prompt_eval_count,
+                            "total_tokens": (
+                                response.prompt_eval_count + response.eval_count
+                                if response.prompt_eval_count is not None
+                                and response.eval_count is not None
+                                else None
+                            ),
+                        },
+                    },
                 )
 
             self.formatter.parse(actions[i], result)
-            results.append(result)
 
-            if generate_logs is not None:
-                # noinspection DuplicatedCode
-                assert isinstance(generate_logs, list)
-                generate_log = GenerateLog()
-                generate_log.prompt = prompts[i]
-                generate_log.backend = f"ollama::{self.model_id!s}"
-                generate_log.date = date
-                generate_log.model_options = model_opts
-                generate_log.model_output = result.value
-                generate_log.extra = {
-                    "format": format,
-                    "thinking": model_opts.get(ModelOption.THINKING, None),
-                    "seed": model_opts.get(ModelOption.SEED, None),
-                }
-                generate_log.action = actions[i]
-                generate_log.result = result
-
-                if error:
-                    generate_log.extra["error"] = error
-                generate_logs.append(generate_log)
+            generate_log = GenerateLog()
+            generate_log.prompt = prompts[i]
+            generate_log.backend = f"ollama::{self.model_id!s}"
+            generate_log.date = date
+            generate_log.model_options = model_opts
+            generate_log.model_output = result.value
+            generate_log.extra = {
+                "format": format,
+                "thinking": model_opts.get(ModelOption.THINKING, None),
+                "seed": model_opts.get(ModelOption.SEED, None),
+            }
+            generate_log.action = actions[i]
+
+            if error:
+                generate_log.extra["error"] = error
+            result._generate_log = generate_log
+
+            results.append(result)
 
         return results
 

mellea/backends/openai.py

@@ -8,7 +8,7 @@
 import json
 from collections.abc import Callable, Coroutine
 from enum import Enum
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, overload
 from urllib.parse import urlparse
 
 import openai
@@ -646,13 +646,14 @@ async def post_processing(
         generate_log.result = mot
         mot._generate_log = generate_log
 
-    def _generate_from_raw(
+    def generate_from_raw(
         self,
         actions: list[Component | CBlock],
+        ctx: Context,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
-        generate_logs: list[GenerateLog] | None = None,
+        tool_calls: bool = False,
     ) -> list[ModelOutputThunk]:
         """Generate using the completions api. Gives the input provided to the model without templating."""
         extra_body = {}
@@ -664,6 +665,10 @@ def _generate_from_raw(
 
             # Some versions (like vllm's version) of the OpenAI API support structured decoding for completions requests.
             extra_body["guided_json"] = format.model_json_schema()
+        if tool_calls:
+            FancyLogger.get_logger().warning(
+                "The completion endpoint does not support tool calling at the moment."
+            )
 
         model_opts = self._simplify_and_merge(model_options, is_chat_context=False)
 
@@ -689,32 +694,35 @@ def _generate_from_raw(
         # Necessary for type checker.
         assert isinstance(completion_response, Completion)
 
-        results = [
-            ModelOutputThunk(
-                value=response.text,
-                meta={"oai_completion_response": response.model_dump()},
-            )
-            for response in completion_response.choices
-        ]
-
-        for i, result in enumerate(results):
-            self.formatter.parse(actions[i], result)
-
-        if generate_logs is not None:
-            assert isinstance(generate_logs, list)
-            date = datetime.datetime.now()
-
-            for i in range(len(prompts)):
-                generate_log = GenerateLog()
-                generate_log.prompt = prompts[i]
-                generate_log.backend = f"openai::{self.model_id!s}"
-                generate_log.model_options = model_opts
-                generate_log.date = date
-                generate_log.model_output = completion_response
-                generate_log.extra = {"seed": model_opts.get("seed", None)}
-                generate_log.action = actions[i]
-                generate_log.result = results[i]
-                generate_logs.append(generate_log)
+        results = []
+        for response, action, prompt in zip(
+            completion_response.choices, actions, prompts
+        ):
+            output = ModelOutputThunk(None)
+            output.value = response.text
+            output._context = None  # There is no context for generate_from_raw for now
+            output._action = action
+            output._model_options = model_opts
+            output._meta = {
+                "oai_completion_response": response.model_dump(),
+                "usage": completion_response.usage.model_dump()
+                if completion_response.usage
+                else None,
+            }
+
+            self.formatter.parse(action, output)
+
+            generate_log = GenerateLog()
+            generate_log.prompt = prompt
+            generate_log.backend = f"openai::{self.model_id!s}"
+            generate_log.model_options = model_opts
+            generate_log.date = datetime.datetime.now()
+            generate_log.model_output = completion_response
+            generate_log.extra = {"seed": model_opts.get("seed", None)}
+            generate_log.action = action
+            output._generate_log = generate_log
+
+            results.append(output)
 
         return results
 

mellea/backends/vllm.py

@@ -393,15 +393,21 @@ async def post_processing(
 
         mot._generate_log = generate_log
 
-    def _generate_from_raw(
+    def generate_from_raw(
         self,
         actions: list[Component | CBlock],
+        ctx: Context,
         *,
         format: type[BaseModelSubclass] | None = None,
         model_options: dict | None = None,
-        generate_logs: list[GenerateLog] | None = None,
+        tool_calls: bool = False,
     ) -> list[ModelOutputThunk]:
         """Generate using the completions api. Gives the input provided to the model without templating."""
+        if tool_calls:
+            FancyLogger.get_logger().warning(
+                "The completion endpoint does not support tool calling at the moment."
+            )
+
         model_options = self._simplify_and_merge(model_options)
 
         prompts = [self.formatter.print(action) for action in actions]
@@ -447,25 +453,21 @@ async def generate_all(prompts):
 
         for i, result in enumerate(results):
             self.formatter.parse(actions[i], result)
-
-        if generate_logs is not None:
-            assert isinstance(generate_logs, list)
             date = datetime.datetime.now()
 
-            for i in range(len(prompts)):
-                generate_log = GenerateLog()
-                generate_log.prompt = prompts[i]
-                generate_log.backend = f"vllm::{self.model_id!s}"
-                generate_log.model_options = model_options
-                generate_log.date = date
-                generate_log.model_output = decoded_results
-                generate_log.extra = {
-                    "format": format,
-                    "seed": model_options.get(ModelOption.SEED, None),
-                }
-                generate_log.action = actions[i]
-                generate_log.result = results[i]
-                generate_logs.append(generate_log)
+            generate_log = GenerateLog()
+            generate_log.prompt = prompts[i]
+            generate_log.backend = f"vllm::{self.model_id!s}"
+            generate_log.model_options = model_options
+            generate_log.date = date
+            generate_log.model_output = decoded_results
+            generate_log.extra = {
+                "format": format,
+                "seed": model_options.get(ModelOption.SEED, None),
+            }
+            generate_log.action = actions[i]
+            generate_log.result = results[i]
+            result._generate_log = generate_log
 
         return results