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

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

Open

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

+275 −2

Collaborator

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)

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