Enhanced Configuration, Token Counting and Unified Output Formatting Schema for swarms Module #770

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Open

harshalmore31 wants to merge 2 commits into kyegomez:master from harshalmore31:new_schema

Contributor

harshalmore31 commented Feb 7, 2025 •

edited by github-actions bot

Loading

This pull request introduces significant enhancements to the swarms module by adding new classes and methods for better handling of token counting, agent configuration, and output formatting. The most important changes include the addition of token counting capabilities, the introduction of new schemas for agent input and swarm configuration, and the creation of a unified output schema.

Token Counting Enhancements:

swarms/schemas/base_schemas.py: Added count_tokens method to ChatMessageInput class to count tokens in message content.
swarms/schemas/base_schemas.py: Added calculate_usage class method to UsageInfo class to calculate token usage for messages and completion.

New Schemas for Agent and Swarm Configuration:

swarms/schemas/base_swarm_schemas.py: Introduced AgentInputConfig class for agent configuration and BaseSwarmSchema class for swarm configuration with detailed validation methods.

Unified Output Schema:

swarms/schemas/output_schemas.py: Added Step, AgentTaskOutput, and OutputSchema classes for representing and formatting the output of agent tasks and swarm executions.

These changes collectively enhance the functionality and robustness of the swarms module, enabling more detailed tracking and configuration of agents and swarms.

📚 Documentation preview 📚: https://swarms--770.org.readthedocs.build/en/770/


          add token counting functionality to message and step schemas

a912702

github-advanced-security bot found potential problems

View reviewed changes

swarms/schemas/base_schemas.py

Comment on lines +118 to +122

    
                      return cls(

                          prompt_tokens=prompt_tokens,

                          completion_tokens=completion_tokens,

                          total_tokens=prompt_tokens + completion_tokens

                      )

Check failure

Code scanning / Pyre

Unexpected keyword Error

Unexpected keyword [28]: Unexpected keyword argument prompt_tokens to call object.__init__.

swarms/schemas/base_swarm_schemas.py

    
            @@ -0,0 +1,152 @@
          
              from typing import Any, Dict, List, Optional, Union

              from pydantic import BaseModel, Field, validator

Check failure

Code scanning / Pyre

Undefined import Error

Undefined import [21]: Could not find a module corresponding to import pydantic.

swarms/schemas/base_swarm_schemas.py

    
              import uuid

              import time

              class AgentInputConfig(BaseModel):

Check failure

Code scanning / Pyre

Undefined or invalid type Error

Undefined or invalid type [11]: Annotation BaseModel is not defined as a type.

swarms/schemas/base_swarm_schemas.py

    
                  Base schema for all swarm types.

                  """

                  id: str = Field(default_factory=lambda: str(uuid.uuid4()))

                  name: str

Check failure

Code scanning / Pyre

Uninitialized attribute Error

Uninitialized attribute [13]: Attribute name is declared in class BaseSwarmSchema to have type str but is never initialized.

swarms/schemas/base_swarm_schemas.py

    
                  """

                  id: str = Field(default_factory=lambda: str(uuid.uuid4()))

                  name: str

                  description: str

Check failure

Code scanning / Pyre

Uninitialized attribute Error

Uninitialized attribute [13]: Attribute description is declared in class BaseSwarmSchema to have type str but is never initialized.

swarms/schemas/base_swarm_schemas.py

    
                  description: str

                  agents: List[AgentInputConfig]  # Using AgentInputConfig

                  max_loops: int = 1

                  swarm_type: str  # e.g., "SequentialWorkflow", "ConcurrentWorkflow", etc.

Check failure

Code scanning / Pyre

Uninitialized attribute Error

Uninitialized attribute [13]: Attribute swarm_type is declared in class BaseSwarmSchema to have type str but is never initialized.

swarms/schemas/output_schemas.py

    
            @@ -0,0 +1,90 @@
          
              from typing import Any, Dict, List, Optional

              from pydantic import BaseModel, Field

Check failure

Code scanning / Pyre

Undefined import Error

Undefined import [21]: Could not find a module corresponding to import pydantic.

swarms/schemas/output_schemas.py

    
              import uuid

              from swarms.utils.litellm_tokenizer import count_tokens

              class Step(BaseModel):

Check failure

Code scanning / Pyre

Undefined or invalid type Error

Undefined or invalid type [11]: Annotation BaseModel is not defined as a type.

swarms/schemas/output_schemas.py

Comment on lines +83 to +89

    
                      output = OutputSchema(

                          swarm_id=swarm_id,

                          swarm_type=swarm_type,

                          task=task,

                          agent_outputs=agent_outputs,

                          swarm_specific_output=swarm_specific_output,

                      )

Check failure

Code scanning / Pyre

Unexpected keyword Error

Unexpected keyword [28]: Unexpected keyword argument swarm_id to call object.__init__.

swarms/schemas/output_schemas.py

    
                          agent_outputs=agent_outputs,

                          swarm_specific_output=swarm_specific_output,

                      )

                      return output.model_dump_json(indent=4)

  No newline at end of file

Check failure

Code scanning / Pyre

Undefined attribute Error

Undefined attribute [16]: OutputSchema has no attribute model_dump_json.


          Updated __init__.py

eb5eca7

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet