autobyteus 1.2.1__py3-none-any.whl → 1.3.0__py3-none-any.whl

autobyteus/agent/workspace/workspace_registry.py

@@ -1,72 +0,0 @@
-"""
-This module provides a central registry for agent workspace types.
-"""
-import logging
-from typing import Dict, Any, Optional, List, TYPE_CHECKING
-from autobyteus.utils.singleton import SingletonMeta
-from .workspace_definition import WorkspaceDefinition
-from .workspace_config import WorkspaceConfig
-
-if TYPE_CHECKING:
-    from .base_workspace import BaseAgentWorkspace
-
-logger = logging.getLogger(__name__)
-
-class WorkspaceRegistry(metaclass=SingletonMeta):
-    """
-    A singleton registry for WorkspaceDefinition objects. Workspaces are
-    typically auto-registered via WorkspaceMeta.
-    """
-    def __init__(self):
-        self._definitions: Dict[str, WorkspaceDefinition] = {}
-        logger.info("Core WorkspaceRegistry initialized.")
-
-    def register(self, definition: WorkspaceDefinition):
-        """Registers a workspace definition."""
-        if not isinstance(definition, WorkspaceDefinition):
-            raise TypeError("Can only register WorkspaceDefinition objects.")
-        if definition.workspace_type_name in self._definitions:
-            logger.warning(f"Overwriting workspace definition for type '{definition.workspace_type_name}'.")
-        self._definitions[definition.workspace_type_name] = definition
-
-    def get_definition(self, workspace_type_name: str) -> Optional[WorkspaceDefinition]:
-        """Retrieves a workspace definition by its unique type name."""
-        return self._definitions.get(workspace_type_name)
-
-    def get_all_definitions(self) -> List[WorkspaceDefinition]:
-        """Returns a list of all registered workspace definitions."""
-        return list(self._definitions.values())
-
-    def create_workspace(self, workspace_type_name: str, config: WorkspaceConfig) -> 'BaseAgentWorkspace':
-        """
-        Creates an instance of a workspace.
-
-        Args:
-            workspace_type_name (str): The unique type name of the workspace to create.
-            config (WorkspaceConfig): The configuration object for the workspace.
-
-        Returns:
-            An instance of a BaseAgentWorkspace subclass.
-            
-        Raises:
-            ValueError: If the type is unknown or parameters are invalid.
-        """
-        definition = self.get_definition(workspace_type_name)
-        if not definition:
-            raise ValueError(f"Unknown workspace type: '{workspace_type_name}'")
-
-        is_valid, errors = definition.config_schema.validate_config(config.to_dict())
-        if not is_valid:
-            error_str = ", ".join(errors)
-            raise ValueError(f"Invalid parameters for workspace type '{workspace_type_name}': {error_str}")
-
-        try:
-            workspace_class = definition.workspace_class
-            instance = workspace_class(config=config)
-            logger.info(f"Successfully created instance of workspace type '{workspace_type_name}'.")
-            return instance
-        except Exception as e:
-            logger.error(f"Failed to instantiate workspace class '{definition.workspace_class.__name__}': {e}", exc_info=True)
-            raise RuntimeError(f"Workspace instantiation failed for type '{workspace_type_name}': {e}") from e
-
-default_workspace_registry = WorkspaceRegistry()

autobyteus/agent_team/bootstrap_steps/agent_team_runtime_queue_initialization_step.py

@@ -1,25 +0,0 @@
-# file: autobyteus/autobyteus/agent_team/bootstrap_steps/agent_team_runtime_queue_initialization_step.py
-import logging
-from typing import TYPE_CHECKING
-
-from autobyteus.agent_team.bootstrap_steps.base_agent_team_bootstrap_step import BaseAgentTeamBootstrapStep
-from autobyteus.agent_team.events.agent_team_input_event_queue_manager import AgentTeamInputEventQueueManager
-
-if TYPE_CHECKING:
-    from autobyteus.agent_team.context.agent_team_context import AgentTeamContext
-    from autobyteus.agent_team.phases.agent_team_phase_manager import AgentTeamPhaseManager
-
-logger = logging.getLogger(__name__)
-
-class AgentTeamRuntimeQueueInitializationStep(BaseAgentTeamBootstrapStep):
-    """Bootstrap step for initializing the agent team's runtime event queues."""
-    async def execute(self, context: 'AgentTeamContext', phase_manager: 'AgentTeamPhaseManager') -> bool:
-        team_id = context.team_id
-        logger.info(f"Team '{team_id}': Executing AgentTeamRuntimeQueueInitializationStep.")
-        try:
-            context.state.input_event_queues = AgentTeamInputEventQueueManager()
-            logger.info(f"Team '{team_id}': AgentTeamInputEventQueueManager initialized.")
-            return True
-        except Exception as e:
-            logger.error(f"Team '{team_id}': Critical failure during queue initialization: {e}", exc_info=True)
-            return False

autobyteus/agent_team/bootstrap_steps/coordinator_prompt_preparation_step.py

@@ -1,85 +0,0 @@
-# file: autobyteus/autobyteus/agent_team/bootstrap_steps/coordinator_prompt_preparation_step.py
-import logging
-from typing import TYPE_CHECKING, List
-
-from autobyteus.agent_team.bootstrap_steps.base_agent_team_bootstrap_step import BaseAgentTeamBootstrapStep
-from autobyteus.agent.context import AgentConfig
-from autobyteus.agent_team.context import AgentTeamConfig
-
-if TYPE_CHECKING:
-    from autobyteus.agent_team.context.agent_team_context import AgentTeamContext
-    from autobyteus.agent_team.phases.agent_team_phase_manager import AgentTeamPhaseManager
-
-logger = logging.getLogger(__name__)
-
-class CoordinatorPromptPreparationStep(BaseAgentTeamBootstrapStep):
-    """
-    Bootstrap step to finalize the coordinator's system prompt by injecting a
-    dynamically generated team manifest into a user-defined prompt template.
-    The user is expected to provide a `system_prompt` in the coordinator's
-    AgentConfig with a `{{team}}` placeholder.
-    """
-    async def execute(self, context: 'AgentTeamContext', phase_manager: 'AgentTeamPhaseManager') -> bool:
-        team_id = context.team_id
-        logger.info(f"Team '{team_id}': Executing CoordinatorPromptPreparationStep.")
-        try:
-            coordinator_node_config_wrapper = context.config.coordinator_node
-            
-            # The coordinator must be an agent with a defined config.
-            if not isinstance(coordinator_node_config_wrapper.node_definition, AgentConfig):
-                logger.error(f"Team '{team_id}': Coordinator node '{coordinator_node_config_wrapper.name}' is not defined by an AgentConfig. Cannot prepare prompt.")
-                return False
-            
-            coordinator_agent_config: AgentConfig = coordinator_node_config_wrapper.node_definition
-            
-            # Start with the user's provided prompt template.
-            prompt_template = coordinator_agent_config.system_prompt
-            if not prompt_template:
-                logger.warning(f"Team '{team_id}': Coordinator '{coordinator_agent_config.name}' has no system_prompt defined. No prompt will be applied.")
-                context.state.prepared_coordinator_prompt = ""
-                return True
-
-            team_manifest = self._generate_team_manifest(context)
-            
-            # Inject the manifest into the template.
-            if "{{team}}" in prompt_template:
-                final_prompt = prompt_template.replace("{{team}}", team_manifest)
-                logger.debug(f"Team '{team_id}': Injected team manifest into coordinator's system prompt.")
-            else:
-                final_prompt = prompt_template
-                logger.warning(f"Team '{team_id}': The coordinator's system prompt does not contain a '{{team}}' placeholder. The team manifest will not be injected.")
-            
-            # Store the finalized prompt in the state for the AgentToolInjectionStep to use.
-            context.state.prepared_coordinator_prompt = final_prompt
-
-            logger.info(f"Team '{team_id}': Coordinator prompt prepared successfully and stored in state.")
-            return True
-        except Exception as e:
-            logger.error(f"Team '{team_id}': Failed to prepare coordinator prompt: {e}", exc_info=True)
-            return False
-
-    def _generate_team_manifest(self, context: 'AgentTeamContext') -> str:
-        """Generates a string manifest of all non-coordinator team members."""
-        prompt_parts: List[str] = []
-        coordinator_node = context.config.coordinator_node
-        member_nodes = {node for node in context.config.nodes if node != coordinator_node}
-
-        if not member_nodes:
-            return "You are working alone. You have no team members to delegate to."
-
-        # Sort for deterministic prompt generation
-        for node in sorted(list(member_nodes), key=lambda n: n.name):
-            node_def = node.node_definition
-            description = "No description available."
-            
-            # --- THE FIX ---
-            # Use the 'description' for an AgentConfig and the 'role' for an AgentTeamConfig (sub-team).
-            if isinstance(node_def, AgentConfig):
-                description = node_def.description
-            elif isinstance(node_def, AgentTeamConfig):
-                # A sub-team's role is its most concise and relevant description for a parent coordinator.
-                description = node_def.role or node_def.description
-            
-            prompt_parts.append(f"- name: {node.name}\n  description: {description}")
-        
-        return "\n".join(prompt_parts)

autobyteus/agent_team/phases/__init__.py

@@ -1,11 +0,0 @@
-# file: autobyteus/autobyteus/agent_team/phases/__init__.py
-"""
-This package contains components for defining and managing agent team operational phases.
-"""
-from autobyteus.agent_team.phases.agent_team_operational_phase import AgentTeamOperationalPhase
-from autobyteus.agent_team.phases.agent_team_phase_manager import AgentTeamPhaseManager
-
-__all__ = [
-    "AgentTeamOperationalPhase",
-    "AgentTeamPhaseManager",
-]

autobyteus/agent_team/phases/agent_team_operational_phase.py

@@ -1,19 +0,0 @@
-# file: autobyteus/autobyteus/agent_team/phases/agent_team_operational_phase.py
-from enum import Enum
-
-class AgentTeamOperationalPhase(str, Enum):
-    """Defines the operational phases of an AgentTeam."""
-    UNINITIALIZED = "uninitialized"
-    BOOTSTRAPPING = "bootstrapping"
-    IDLE = "idle"
-    PROCESSING = "processing"
-    SHUTTING_DOWN = "shutting_down"
-    SHUTDOWN_COMPLETE = "shutdown_complete"
-    ERROR = "error"
-
-    def is_terminal(self) -> bool:
-        """Checks if the phase is a terminal state."""
-        return self in [AgentTeamOperationalPhase.SHUTDOWN_COMPLETE, AgentTeamOperationalPhase.ERROR]
-
-    def __str__(self) -> str:
-        return self.value

autobyteus/agent_team/phases/agent_team_phase_manager.py

@@ -1,48 +0,0 @@
-# file: autobyteus/autobyteus/agent_team/phases/agent_team_phase_manager.py
-import logging
-from typing import TYPE_CHECKING, Optional
-
-from autobyteus.agent_team.phases.agent_team_operational_phase import AgentTeamOperationalPhase
-
-if TYPE_CHECKING:
-    from autobyteus.agent_team.context.agent_team_context import AgentTeamContext
-    from autobyteus.agent_team.streaming.agent_team_event_notifier import AgentTeamExternalEventNotifier
-
-logger = logging.getLogger(__name__)
-
-class AgentTeamPhaseManager:
-    """Manages the operational phase of an agent team."""
-    def __init__(self, context: 'AgentTeamContext', notifier: 'AgentTeamExternalEventNotifier'):
-        self.context = context
-        self.notifier = notifier
-        self.context.state.current_phase = AgentTeamOperationalPhase.UNINITIALIZED
-        logger.debug(f"AgentTeamPhaseManager initialized for team '{context.team_id}'.")
-
-    async def _transition_phase(self, new_phase: AgentTeamOperationalPhase, extra_data: Optional[dict] = None):
-        old_phase = self.context.state.current_phase
-        if old_phase == new_phase:
-            return
-        logger.info(f"Team '{self.context.team_id}' transitioning from {old_phase.value} to {new_phase.value}.")
-        self.context.state.current_phase = new_phase
-        self.notifier.notify_phase_change(new_phase, old_phase, extra_data)
-
-    async def notify_bootstrapping_started(self):
-        await self._transition_phase(AgentTeamOperationalPhase.BOOTSTRAPPING)
-
-    async def notify_initialization_complete(self):
-        await self._transition_phase(AgentTeamOperationalPhase.IDLE)
-        
-    async def notify_processing_started(self):
-        await self._transition_phase(AgentTeamOperationalPhase.PROCESSING)
-
-    async def notify_processing_complete_and_idle(self):
-        await self._transition_phase(AgentTeamOperationalPhase.IDLE)
-
-    async def notify_error_occurred(self, error_message: str, error_details: Optional[str] = None):
-        await self._transition_phase(AgentTeamOperationalPhase.ERROR, {"error_message": error_message, "error_details": error_details})
-
-    async def notify_shutdown_initiated(self):
-        await self._transition_phase(AgentTeamOperationalPhase.SHUTTING_DOWN)
-
-    async def notify_final_shutdown_complete(self):
-        await self._transition_phase(AgentTeamOperationalPhase.SHUTDOWN_COMPLETE)

autobyteus/llm/api/bedrock_llm.py

@@ -1,92 +0,0 @@
-from typing import Dict, Optional, List, AsyncGenerator
-import boto3
-import json
-import os
-from botocore.exceptions import ClientError
-from autobyteus.llm.models import LLMModel
-from autobyteus.llm.base_llm import BaseLLM
-from autobyteus.llm.utils.llm_config import LLMConfig
-from autobyteus.llm.utils.messages import MessageRole, Message
-from autobyteus.llm.utils.token_usage import TokenUsage
-from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
-from autobyteus.llm.user_message import LLMUserMessage
-
-class BedrockLLM(BaseLLM):
-    def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
-        if model is None:
-            model = LLMModel.BEDROCK_CLAUDE_3_5_SONNET_API
-        if llm_config is None:
-            llm_config = LLMConfig()
-            
-        super().__init__(model=model, llm_config=llm_config)
-        self.client = self.initialize()
-    
-    @classmethod
-    def initialize(cls):
-        aws_access_key = os.environ.get("AWS_ACCESS_KEY_ID")
-        aws_secret_key = os.environ.get("AWS_SECRET_ACCESS_KEY")
-        region = os.environ.get("AWS_REGION", "us-east-1")
-
-        if not (aws_access_key and aws_secret_key):
-            raise ValueError(
-                "AWS credentials not found. Please set AWS_ACCESS_KEY_ID and "
-                "AWS_SECRET_ACCESS_KEY environment variables."
-            )
-
-        try:
-            return boto3.client(
-                service_name='bedrock-runtime',
-                region_name=region,
-                aws_access_key_id=aws_access_key,
-                aws_secret_access_key=aws_secret_key
-            )
-        except Exception as e:
-            raise ValueError(f"Failed to initialize Bedrock client: {str(e)}")
-    
-    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
-        self.add_user_message(user_message)
-        
-        # NOTE: This implementation does not yet support multimodal inputs for Bedrock.
-        # It will only send the text content.
-        
-        request_body = json.dumps({
-            "anthropic_version": "bedrock-2023-05-31",
-            "max_tokens": 1000,
-            "temperature": 0,
-            "messages": [msg.to_dict() for msg in self.messages if msg.role != MessageRole.SYSTEM],
-            "system": self.system_message if self.system_message else ""
-        })
-
-        try:
-            response = self.client.invoke_model(
-                modelId=self.model.value,
-                body=request_body
-            )
-            response_body = json.loads(response['body'].read())
-            assistant_message = response_body['content'][0]['text']
-            self.add_assistant_message(assistant_message)
-            
-            token_usage = TokenUsage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0
-            )
-            
-            return CompleteResponse(
-                content=assistant_message,
-                usage=token_usage
-            )
-        except ClientError as e:
-            error_code = e.response['Error']['Code']
-            error_message = e.response['Error']['Message']
-            raise ValueError(f"Bedrock API error: {error_code} - {error_message}")
-        except Exception as e:
-            raise ValueError(f"Error in Bedrock API call: {str(e)}")
-
-    async def _stream_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> AsyncGenerator[ChunkResponse, None]:
-        # Placeholder for future implementation
-        response = await self._send_user_message_to_llm(user_message, **kwargs)
-        yield ChunkResponse(content=response.content, is_complete=True, usage=response.usage)
-
-    async def cleanup(self):
-        await super().cleanup()

autobyteus/llm/api/groq_llm.py

@@ -1,94 +0,0 @@
-from typing import Dict, Optional, List, AsyncGenerator
-import logging
-import os
-from autobyteus.llm.models import LLMModel
-from autobyteus.llm.base_llm import BaseLLM
-from autobyteus.llm.utils.llm_config import LLMConfig
-from autobyteus.llm.utils.messages import MessageRole, Message
-from autobyteus.llm.utils.token_usage import TokenUsage
-from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
-from autobyteus.llm.user_message import LLMUserMessage
-
-logger = logging.getLogger(__name__)
-
-class GroqLLM(BaseLLM):
-    def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
-        # Provide defaults if not specified
-        if model is None:
-            model = LLMModel.LLAMA_3_1_70B_VERSATILE_API
-        if llm_config is None:
-            llm_config = LLMConfig()
-            
-        super().__init__(model=model, llm_config=llm_config)
-        self.client = self.initialize()
-
-    @classmethod
-    def initialize(cls):
-        groq_api_key = os.environ.get("GROQ_API_KEY")
-        if not groq_api_key:
-            raise ValueError(
-                "GROQ_API_KEY environment variable is not set. "
-                "Please set this variable in your environment."
-            )
-        try:
-            # Initialize Groq client here
-            # Placeholder for actual initialization
-            return "GroqClientInitialized"
-        except Exception as e:
-            raise ValueError(f"Failed to initialize Groq client: {str(e)}")
-    
-    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
-        self.add_user_message(user_message)
-        try:
-            # Placeholder for sending message to Groq API
-            assistant_message = "Response from Groq API"
-            self.add_assistant_message(assistant_message)
-            
-            token_usage = TokenUsage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0
-            )
-            
-            return CompleteResponse(
-                content=assistant_message,
-                usage=token_usage
-            )
-        except Exception as e:
-            logger.error(f"Error in Groq API call: {str(e)}")
-            raise ValueError(f"Error in Groq API call: {str(e)}")
-    
-    async def _stream_user_message_to_llm(
-        self, user_message: LLMUserMessage, **kwargs
-    ) -> AsyncGenerator[ChunkResponse, None]:
-        self.add_user_message(user_message)
-        complete_response = ""
-        try:
-            # Placeholder for streaming from Groq API
-            tokens = ["Response ", "streamed ", "from ", "Groq ", "API."]
-            for token in tokens:
-                complete_response += token
-                yield ChunkResponse(
-                    content=token,
-                    is_complete=False
-                )
-            
-            token_usage = TokenUsage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0
-            )
-            
-            yield ChunkResponse(
-                content="",
-                is_complete=True,
-                usage=token_usage
-            )
-            
-            self.add_assistant_message(complete_response)
-        except Exception as e:
-            logger.error(f"Error in Groq API streaming: {str(e)}")
-            raise ValueError(f"Error in Groq API streaming: {str(e)}")
-    
-    async def cleanup(self):
-        await super().cleanup()

autobyteus/llm/api/nvidia_llm.py

@@ -1,108 +0,0 @@
-import logging
-from typing import Dict, Optional, List, AsyncGenerator
-from openai import OpenAI
-import os
-from autobyteus.llm.models import LLMModel
-from autobyteus.llm.base_llm import BaseLLM
-from autobyteus.llm.utils.llm_config import LLMConfig
-from autobyteus.llm.utils.messages import MessageRole, Message
-from autobyteus.llm.utils.token_usage import TokenUsage
-from autobyteus.llm.utils.response_types import CompleteResponse, ChunkResponse
-from autobyteus.llm.user_message import LLMUserMessage
-
-logger = logging.getLogger(__name__)
-
-class NvidiaLLM(BaseLLM):
-    def __init__(self, model: LLMModel = None, llm_config: LLMConfig = None):
-        # Provide defaults if not specified
-        if model is None:
-            model = LLMModel.NVIDIA_LLAMA_3_1_NEMOTRON_70B_INSTRUCT_API
-        if llm_config is None:
-            llm_config = LLMConfig()
-            
-        super().__init__(model=model, llm_config=llm_config)
-        self.client = self.initialize()
-    
-    @classmethod
-    def initialize(cls):
-        nvidia_api_key = os.environ.get("NVIDIA_API_KEY")
-        if not nvidia_api_key:
-            raise ValueError(
-                "NVIDIA_API_KEY environment variable is not set. "
-                "Please set this variable in your environment."
-            )
-        try:
-            return OpenAI(
-                base_url="https://integrate.api.nvidia.com/v1",
-                api_key=nvidia_api_key
-            )
-        except Exception as e:
-            raise ValueError(f"Failed to initialize Nvidia client: {str(e)}")
-    
-    async def _send_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> CompleteResponse:
-        self.add_user_message(user_message)
-        try:
-            completion = self.client.chat.completions.create(
-                model=self.model.value,
-                messages=[msg.to_dict() for msg in self.messages],
-                temperature=0,
-                top_p=1,
-                max_tokens=1024,
-                stream=False
-            )
-            assistant_message = completion.choices[0].message.content
-            self.add_assistant_message(assistant_message)
-            
-            token_usage = TokenUsage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0
-            )
-            
-            return CompleteResponse(
-                content=assistant_message,
-                usage=token_usage
-            )
-        except Exception as e:
-            raise ValueError(f"Error in Nvidia API call: {str(e)}")
-    
-    async def _stream_user_message_to_llm(self, user_message: LLMUserMessage, **kwargs) -> AsyncGenerator[ChunkResponse, None]:
-        self.add_user_message(user_message)
-        complete_response = ""
-        try:
-            completion = self.client.chat.completions.create(
-                model=self.model.value,
-                messages=[msg.to_dict() for msg in self.messages],
-                temperature=0,
-                top_p=1,
-                max_tokens=1024,
-                stream=True
-            )
-            
-            for chunk in completion:
-                if chunk.choices[0].delta.content is not None:
-                    token = chunk.choices[0].delta.content
-                    complete_response += token
-                    yield ChunkResponse(
-                        content=token,
-                        is_complete=False
-                    )
-            
-            token_usage = TokenUsage(
-                prompt_tokens=0,
-                completion_tokens=0,
-                total_tokens=0
-            )
-            
-            yield ChunkResponse(
-                content="",
-                is_complete=True,
-                usage=token_usage
-            )
-
-            self.add_assistant_message(complete_response)
-        except Exception as e:
-            raise ValueError(f"Error in Nvidia API streaming call: {str(e)}")
-    
-    async def cleanup(self):
-        await super().cleanup()

autobyteus/llm/utils/token_pricing_config.py

@@ -1,87 +0,0 @@
-
-from dataclasses import dataclass
-from autobyteus.llm.models import LLMModel
-
-@dataclass
-class TokenPricingConfig:
-    """
-    Represents the pricing configuration for input and output tokens.
-    Prices are in USD per million tokens.
-    """
-    input_token_pricing: float  # USD per million tokens
-    output_token_pricing: float  # USD per million tokens
-
-class TokenPricingConfigRegistry:
-    """
-    Registry for token pricing configurations of different LLM models.
-    Uses model names (enum values) as keys for simplicity.
-    """
-    _prices = {
-        # ChatGPT models
-        LLMModel.GPT_4o_API.value: TokenPricingConfig(2.50, 10.00),
-        LLMModel.o1_API.value: TokenPricingConfig(15.00, 60.00),
-        LLMModel.o1_MINI_API.value: TokenPricingConfig(3.00, 12.00),
-        LLMModel.CHATGPT_4O_LATEST_API.value: TokenPricingConfig(2.50, 10.00),
-        LLMModel.GPT_3_5_TURBO_API.value: TokenPricingConfig(1.50, 2.00),
-
-        # Mistral models
-        LLMModel.MISTRAL_SMALL_API.value: TokenPricingConfig(0.20, 0.60),
-        LLMModel.MISTRAL_MEDIUM_API.value: TokenPricingConfig(0.20, 0.60),
-        LLMModel.MISTRAL_LARGE_API.value: TokenPricingConfig(2.00, 6.00),
-
-        # Groq models
-        LLMModel.GEMMA_2_9B_IT_API.value: TokenPricingConfig(0.0000005, 0.0000005),
-        LLMModel.GEMMA_7B_IT_API.value: TokenPricingConfig(0.0000007, 0.0000007),
-        LLMModel.LLAMA_3_1_405B_REASONING_API.value: TokenPricingConfig(0.00002, 0.00002),
-        LLMModel.LLAMA_3_1_70B_VERSATILE_API.value: TokenPricingConfig(0.00001, 0.00001),
-        LLMModel.LLAMA_3_1_8B_INSTANT_API.value: TokenPricingConfig(0.0000005, 0.0000005),
-        LLMModel.LLAMA3_70B_8192_API.value: TokenPricingConfig(0.00001, 0.00001),
-        LLMModel.LLAMA3_8B_8192_API.value: TokenPricingConfig(0.0000005, 0.0000005),
-        LLMModel.MIXTRAL_8X7B_32768_API.value: TokenPricingConfig(0.000001, 0.000001),
-
-        # Gemini models
-        LLMModel.GEMINI_1_0_PRO_API.value: TokenPricingConfig(0.00000025, 0.0000005),
-        LLMModel.GEMINI_1_5_PRO_API.value: TokenPricingConfig(0.0000005, 0.000001),
-        LLMModel.GEMINI_1_5_PRO_EXPERIMENTAL_API.value: TokenPricingConfig(0.0000005, 0.000001),
-        LLMModel.GEMINI_1_5_FLASH_API.value: TokenPricingConfig(0.0000001, 0.0000002),
-        LLMModel.GEMMA_2_2B_API.value: TokenPricingConfig(0.0000001, 0.0000002),
-        LLMModel.GEMMA_2_9B_API.value: TokenPricingConfig(0.0000001, 0.0000002),
-        LLMModel.GEMMA_2_27B_API.value: TokenPricingConfig(0.0000001, 0.0000002),
-
-        # Claude models
-        LLMModel.CLAUDE_3_HAIKU_API.value: TokenPricingConfig(0.25, 1.25),
-        LLMModel.CLAUDE_3_OPUS_API.value: TokenPricingConfig(15.00, 75.00),
-        LLMModel.CLAUDE_3_5_SONNET_API.value: TokenPricingConfig(3.00, 15.00),
-        LLMModel.CLAUDE_3_SONNET_API.value: TokenPricingConfig(3.00, 15.00),
-        LLMModel.BEDROCK_CLAUDE_3_5_SONNET_API.value: TokenPricingConfig(3.00, 15.00),
-
-        # Perplexity models
-        LLMModel.LLAMA_3_1_SONAR_LARGE_128K_ONLINE_API.value: TokenPricingConfig(0.000001, 0.000001),
-        LLMModel.LLAMA_3_1_SONAR_SMALL_128K_ONLINE_API.value: TokenPricingConfig(0.0000005, 0.0000005),
-        LLMModel.LLAMA_3_1_SONAR_LARGE_128K_CHAT_API.value: TokenPricingConfig(0.000001, 0.000001),
-        LLMModel.LLAMA_3_1_SONAR_SMALL_128K_CHAT_API.value: TokenPricingConfig(0.0000005, 0.0000005),
-        LLMModel.LLAMA_3_1_8B_INSTRUCT_API.value: TokenPricingConfig(0.0000005, 0.0000005),
-        LLMModel.LLAMA_3_1_70B_INSTRUCT_API.value: TokenPricingConfig(0.00001, 0.00001),
-        LLMModel.GEMMA_2_27B_IT_API.value: TokenPricingConfig(0.000001, 0.000001),
-        LLMModel.NEMOTRON_4_340B_INSTRUCT_API.value: TokenPricingConfig(0.00002, 0.00002),
-        LLMModel.MIXTRAL_8X7B_INSTRUCT_API.value: TokenPricingConfig(0.000001, 0.000001),
-
-        # DeepSeek models
-        LLMModel.DEEPSEEK_CHAT_API.value: TokenPricingConfig(0.14, 0.28),
-
-        # NVIDIA models
-        LLMModel.NVIDIA_LLAMA_3_1_NEMOTRON_70B_INSTRUCT_API.value: TokenPricingConfig(0.00002, 0.00002),
-    }
-
-    @classmethod
-    def get_pricing(cls, model_name: str) -> TokenPricingConfig:
-        """
-        Get the token pricing configuration for a given model name.
-        
-        Args:
-            model_name (str): The name of the model to get pricing for.
-            
-        Returns:
-            TokenPricingConfig: The token pricing configuration for the model.
-        """
-        return cls._prices.get(model_name, TokenPricingConfig(0, 0))