[FEAT][GenerationOutputMetadata]

playground/agents/agent_with_basemodel_output_type.py

@@ -12,6 +12,7 @@ class Schema(BaseModel):
         ..., title="List of courses the person is taking"
     )
 
+
 # Convert the schema to a JSON string
 tool_schema = Schema(
     name="Tool Name",
@@ -39,12 +40,12 @@ class Schema(BaseModel):
     verbose=True,
     interactive=True,
     # Set the output type to the tool schema which is a BaseModel
-    output_type=tool_schema, # or dict, or str
+    output_type=tool_schema,  # or dict, or str
     metadata_output_type="json",
     # List of schemas that the agent can handle
-    list_tool_schemas = [tool_schema],
-    function_calling_format_type = "OpenAI",
-    function_calling_type = "json" # or soon yaml
+    list_tool_schemas=[tool_schema],
+    function_calling_format_type="OpenAI",
+    function_calling_type="json",  # or soon yaml
 )
 
 # Run the agent to generate the person's information

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "swarms"
-version = "4.8.6"
+version = "4.8.7"
 description = "Swarms - Pytorch"
 license = "MIT"
 authors = ["Kye Gomez <[email protected]>"]

swarms/models/popular_llms.py

@@ -7,10 +7,11 @@
 from langchain.llms.anthropic import Anthropic
 from langchain.llms.cohere import Cohere
 from langchain.llms.mosaicml import MosaicML
-from langchain.llms.openai import OpenAI #, OpenAIChat, AzureOpenAI
+from langchain.llms.openai import OpenAI  # , OpenAIChat, AzureOpenAI
 from langchain_community.llms.octoai_endpoint import OctoAIEndpoint
 from langchain.llms.replicate import Replicate
 
+
 class AnthropicChat(Anthropic):
     def __call__(self, *args, **kwargs):
         return self.invoke(*args, **kwargs)
@@ -44,7 +45,7 @@ def __call__(self, *args, **kwargs):
 class OpenAIChatLLM(OpenAIChat):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        
+
     def __call__(self, *args, **kwargs):
         return self.invoke(*args, **kwargs)
 

swarms/structs/agent.py

@@ -323,10 +323,7 @@ def __init__(
 
         # If the stopping function is provided then set the stopping condition to the stopping function
         self.short_memory = Conversation(
-            system_prompt=system_prompt,
-            time_enabled=True,
-            *args,
-            **kwargs
+            system_prompt=system_prompt, time_enabled=True, *args, **kwargs
         )
 
         # If the docs exist then ingest the docs
@@ -392,7 +389,7 @@ def __init__(
         if self.tool_schema is not None:
             logger.info("Tool schema provided")
             tool_schema_str = self.tool_schema_to_str(self.tool_schema)
-            
+
             print(tool_schema_str)
 
             # Add to the short memory
@@ -467,6 +464,48 @@ def add_task_to_memory(self, task: str):
         except Exception as error:
             print(colored(f"Error adding task to memory: {error}", "red"))
 
+    # ############## TOKENIZER FUNCTIONS ##############
+    def count_tokens(self, text: str) -> int:
+        """Count the number of tokens in the text."""
+        return self.tokenizer.len(text)
+
+    def tokens_per_second(self, text: str) -> float:
+        """
+        Calculates the number of tokens processed per second.
+
+        Args:
+            text (str): The input text to count tokens from.
+
+        Returns:
+            float: The number of tokens processed per second.
+        """
+        import time
+
+        start_time = time.time()
+        tokens = self.count_tokens(text)
+        end_time = time.time()
+        elapsed_time = end_time - start_time
+        return tokens / elapsed_time
+
+    def time_to_generate(self, text: str) -> float:
+        """
+        Calculates the time taken to generate the output.
+
+        Args:
+            text (str): The input text to generate output from.
+
+        Returns:
+            float: The time taken to generate the output.
+        """
+        import time
+
+        start_time = time.time()
+        self.llm(text)
+        end_time = time.time()
+        return end_time - start_time
+
+    # ############## TOKENIZER FUNCTIONS ##############
+
     def add_message_to_memory(self, message: str):
         """Add the message to the memory"""
         try:

swarms/structs/base_structure.py

@@ -80,7 +80,6 @@ def __init__(
         self.save_artifact_path = save_artifact_path
         self.save_metadata_path = save_metadata_path
         self.save_error_path = save_error_path
-
 
     def run(self, *args, **kwargs):
         """Run the structure."""

swarms/structs/schemas.py

@@ -164,3 +164,91 @@ class ManySteps(BaseModel):
         [],
         description="A list of task steps.",
     )
+
+
+class GenerationOutputMetadata(BaseModel):
+    num_of_tokens: int = Field(
+        ...,
+        description="The number of tokens generated.",
+        examples=[7894],
+    )
+    estimated_cost: str = Field(
+        ...,
+        description="The estimated cost of the generation.",
+        examples=["0,24$"],
+    )
+    time_to_generate: str = Field(
+        ...,
+        description="The time taken to generate the output.",
+        examples=["1.2s"],
+    )
+    tokens_per_second: int = Field(
+        ...,
+        description="The number of tokens generated per second.",
+        examples=[657],
+    )
+    model_name: str = Field(
+        ...,
+        description="The model used to generate the output.",
+        examples=["gpt-3.5-turbo"],
+    )
+    max_tokens: int = Field(
+        ...,
+        description="The maximum number of tokens allowed to generate.",
+        examples=[2048],
+    )
+    temperature: float = Field(
+        ...,
+        description="The temperature used for generation.",
+        examples=[0.7],
+    )
+    top_p: float = Field(
+        ...,
+        description="The top p value used for generation.",
+        examples=[0.9],
+    )
+    frequency_penalty: float = Field(
+        ...,
+        description="The frequency penalty used for generation.",
+        examples=[0.0],
+    )
+    presence_penalty: float = Field(
+        ...,
+        description="The presence penalty used for generation.",
+        examples=[0.0],
+    )
+    stop_sequence: str | None = Field(
+        None,
+        description="The sequence used to stop the generation.",
+        examples=["<stop_sequence>"],
+    )
+    model_type: str = Field(
+        ...,
+        description="The type of model used for generation.",
+        examples=["text"],
+    )
+    model_version: str = Field(
+        ...,
+        description="The version of the model used for generation.",
+        examples=["1.0.0"],
+    )
+    model_description: str = Field(
+        ...,
+        description="The description of the model used for generation.",
+        examples=["A model that generates text."],
+    )
+    model_author: str = Field(
+        ...,
+        description="The author of the model used for generation.",
+        examples=["John Doe"],
+    )
+    n: int = Field(
+        ...,
+        description="The number of outputs generated.",
+        examples=[1],
+    )
+    n_best: int = Field(
+        ...,
+        description="The number of best outputs generated.",
+        examples=[1],
+    )

swarms/tools/pydantic_to_json.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, List
+from typing import Any, List
 
 from docstring_parser import parse
 from pydantic import BaseModel
@@ -72,7 +72,7 @@ def pydantic_to_functions(
 
 
 def multi_pydantic_to_functions(
-    pydantic_types: List[BaseModel] = None
+    pydantic_types: List[BaseModel] = None,
 ) -> dict[str, Any]:
     """
     Converts multiple Pydantic types to a dictionary of functions.