Dynamic models, o-models, nebius support (#32)

Author: maisiukartyom

.github/workflows/run-unit-tests.yaml

@@ -19,8 +19,7 @@ jobs:
         echo OPENAI_API_KEY=${{ secrets.OPENAI_API_KEY }} >> .env
         echo LAMOOM_API_URI=${{ secrets.LAMOOM_API_URI }} >> .env
         echo LAMOOM_API_TOKEN=${{ secrets.LAMOOM_API_TOKEN }} >> .env
-        echo FLOW_PROMPT_API_URI=${{ secrets.FLOW_PROMPT_API_URI }} >> .env
-        echo FLOW_PROMPT_API_TOKEN=${{ secrets.FLOW_PROMPT_API_TOKEN }} >> .env
+        cat .env
 
     - name: Install dependencies
       run: |
@@ -41,3 +40,9 @@ jobs:
     - name: Run tests with pytest
       run: |
         poetry run make test
+
+    - name: Publish package
+      env:
+        PYPI_API_KEY: ${{ secrets.PYPI_API_KEY }}
+      run: |
+        poetry run make publish-release

Makefile

@@ -62,6 +62,7 @@ publish-test-prerelease:
 
 
 publish-release:
+	poetry config pypi-token.pypi "$(PYPI_API_KEY)"
 	poetry version patch
 	poetry build
 	poetry publish

README.md

@@ -35,12 +35,8 @@ os.setenv('OPENAI_API_KEY', 'your_key_here')
 
 # add Azure Keys
 os.setenv('AZURE_KEYS', '{"name_realm":{"url": "https://baseurl.azure.com/","key": "secret"}}')
-
-# Other env Variables;
-# CLAUDE_API_KEY
-# GEMINI_API_KEY
-# OPENAI_ORG
-# NEBIUS_KEY
+# or creating flow_prompt obj
+Lamoom(azure_keys={"realm_name":{"url": "https://baseurl.azure.com/", "key": "your_secret"}})
 ```
 
 ### Model Agnostic:
@@ -49,40 +45,25 @@ Mix models easily, and districute the load across models. The system will automa
 - Gemini
 - OpenAI (w/ Azure OpenAI models)
 - Nebius with (Llama, DeepSeek, Mistral, Mixtral, dolphin, Qwen and others)
-```
-from lamoom import LamoomModelProviders
-
-def_behaviour = behaviour.AIModelsBehaviour(attempts=[
-    AttemptToCall(provider='openai', model='gpt-4o', weight=100),
-    AttemptToCall(provider='azure', realm='useast-1', deployment_id='gpt-4o' weight=100),
-    AttemptToCall(provider='azure', realm='useast-2', deployment_id='gpt-4o' weight=100),
-    AttemptToCall(provider=LamoomModelProviders.anthropic, model='claude-3-5-sonnet-20240620', weight=100
-    ),
-    AttemptToCall(provider=LamoomModelProviders.gemini, model='gemini-1.5-pro', weight=100
-    ),
-    AttemptToCall(provider=LamoomModelProviders.nebius, model='deepseek-ai/DeepSeek-R1', weight=100
-    )
-])
-
-response_llm = client.call(agent.id, context, def_behaviour)
-```
 
-### Add Behavious:
-- use OPENAI_BEHAVIOR
-- or add your own Behaviour, you can set max count of attempts, if you have different AI Models, if the first attempt will fail because of retryable error, the second will be called, based on the weights.
-```
-from lamoom import OPENAI_GPT4_0125_PREVIEW_BEHAVIOUR
-behaviour = OPENAI_GPT4_0125_PREVIEW_BEHAVIOUR
-```
-or:
+Model string format is the following for Claude, Gemini, OpenAI, Nebius:
+`"{model_provider}/{model_name}"`
+For Azure models format is the following:
+`"azure/{realm}/{model_name}"`
+
+```python
+response_llm = client.call(agent.id, context, model = "openai/gpt-4o")
+response_llm = client.call(agent.id, context, model = "azure/useast/gpt-4o")
 ```
-from lamoom import behaviour
-behaviour = behaviour.AIModelsBehaviour(
-    attempts=[
-        AttemptToCall(provider='azure', realm='useast-1', deployment_id='gpt-4o' weight=100),
-        AttemptToCall(provider='azure', realm='useast-2', deployment_id='gpt-4o' weight=100),
-    ]
-)
+
+### Lamoom Keys
+Obtain an API token from Flow Prompt and add it:
+
+```python
+# As an environment variable:
+os.setenv('LAMOOM_API_TOKEN', 'your_token_here')
+# Via code: 
+Lamoom(api_token='your_api_token')
 ```
 
 ## Usage Examples:
@@ -100,7 +81,7 @@ prompt.add("You're {name}. Say Hello and ask what's their name.", role="system")
 # Call AI model with Lamoom
 context = {"name": "John Doe"}
 # test_data -  optional parameter used for generating tests
-response = client.call(prompt.id, context, behavior, test_data={
+response = client.call(prompt.id, context, "openai/gpt-4o", test_data={
     'ideal_answer': "Hello, I'm John Doe. What's your name?", 
     'behavior_name': "gemini"
     }

lamoom/__init__.py

@@ -4,14 +4,6 @@
 from lamoom.ai_models import behaviour
 from lamoom.prompt.prompt import Prompt
 from lamoom.prompt.prompt import Prompt as PipePrompt
-from lamoom.ai_models.openai.behaviours import (
-    OPENAI_GPT4_0125_PREVIEW_BEHAVIOUR,
-    OPENAI_GPT4_1106_PREVIEW_BEHAVIOUR,
-    OPENAI_GPT4_1106_VISION_PREVIEW_BEHAVIOUR,
-    OPENAI_GPT4_BEHAVIOUR,
-    OPENAI_GPT4_32K_BEHAVIOUR,
-    OPENAI_GPT3_5_TURBO_0125_BEHAVIOUR,
-)
 from lamoom.ai_models.attempt_to_call import AttemptToCall
 from lamoom.ai_models.openai.openai_models import (
     C_128K,
@@ -22,7 +14,6 @@
 )
 from lamoom.ai_models.openai.azure_models import AzureAIModel
 from lamoom.ai_models.claude.claude_model import ClaudeAIModel
-from lamoom.ai_models.gemini.gemini_model import GeminiAIModel
 from lamoom.responses import AIResponse
 from lamoom.ai_models.openai.responses import OpenAIResponse
 from lamoom.ai_models.behaviour import AIModelsBehaviour, PromptAttempts

lamoom/ai_models/ai_model.py

@@ -10,44 +10,24 @@
 class AI_MODELS_PROVIDER(Enum):
     OPENAI = "openai"
     AZURE = "azure"
-    CLAUDE = ("claude",)
+    CLAUDE = "claude"
     GEMINI = "gemini"
+    NEBIUS = "nebius"
 
 
 @dataclass(kw_only=True)
 class AIModel:
-    max_tokens: int
     tiktoken_encoding: t.Optional[str] = "cl100k_base"
     provider: AI_MODELS_PROVIDER = None
     support_functions: bool = False
-    _price_per_prompt_1k_tokens: Decimal = None
-    _price_per_sample_1k_tokens: Decimal = None
 
     @property
     def name(self) -> str:
         return "undefined_aimodel"
 
-    @property
-    def price_per_prompt_1k_tokens(self) -> Decimal:
-        return self._price_per_prompt_1k_tokens
-
     def _decimal(self, value) -> Decimal:
         return Decimal(value).quantize(Decimal(".00001"))
 
-    def get_prompt_price(self, count_tokens: int) -> Decimal:
-        return self._decimal(
-            self.price_per_prompt_1k_tokens * Decimal(count_tokens) / 1000
-        )
-
-    def get_sample_price(self, prompt_sample, count_tokens: int) -> Decimal:
-        return self._decimal(
-            self.price_per_sample_1k_tokens * Decimal(count_tokens) / 1000
-        )
-
-    @property
-    def price_per_sample_1k_tokens(self) -> Decimal:
-        return self._price_per_sample_1k_tokens
-
     def get_params(self) -> t.Dict[str, t.Any]:
         return {}
 

lamoom/ai_models/behaviour.py

@@ -12,51 +12,28 @@
 
 @dataclass
 class AIModelsBehaviour:
-    # if you have mutiple AI Models, you can distribute the load across them.
-    # If you wish to use as a fallback attempt a model which is not in the list, you can use fallback_attempt
-    attempts: list[AttemptToCall]
-    fallback_attempt: AttemptToCall = None
+    attempt: AttemptToCall
+    fallback_attempts: list[AttemptToCall] = None
 
 
 @dataclass
 class PromptAttempts:
     ai_models_behaviour: AIModelsBehaviour
-    count_of_retries: t.Optional[int] = None
-    count: int = 0
     current_attempt: AttemptToCall = None
 
-    def __post_init__(self):
-        if self.count_of_retries is None:
-            self.count_of_retries = len(self.ai_models_behaviour.attempts) + int(
-                bool(self.ai_models_behaviour.fallback_attempt)
-            )
-
-    def initialize_attempt(self, flag_increase_count: bool = True):
-        if self.count > self.count_of_retries:
-            raise BehaviourIsNotDefined(
-                f"Count of retries {self.count_of_retries} exceeded {self.count}"
-            )
-        if (
-            self.count == self.count_of_retries
-            and self.ai_models_behaviour.fallback_attempt
-        ):
-            self.current_attempt = self.ai_models_behaviour.fallback_attempt
+    def initialize_attempt(self):
+        if self.current_attempt is None:
+            self.current_attempt = self.ai_models_behaviour.attempt
+            self.fallback_index = 0  # Start fallback index at 0
             return self.current_attempt
-        sum_weight = sum(
-            [attempt.weight for attempt in self.ai_models_behaviour.attempts]
-        )
-        random_weight = random.randint(0, sum_weight)
-        for attempt in self.ai_models_behaviour.attempts:
-            random_weight -= attempt.weight
-            if random_weight <= 0:
-                if flag_increase_count:
-                    self.count += 1
-                attempt.attempt_number = self.count
-                self.current_attempt = copy(attempt)
+        elif self.ai_models_behaviour.fallback_attempts:
+            if self.fallback_index < len(self.ai_models_behaviour.fallback_attempts):
+                self.current_attempt = self.ai_models_behaviour.fallback_attempts[self.fallback_index]
+                self.fallback_index += 1
                 return self.current_attempt
-        raise BehaviourIsNotDefined(
-            f"Count of retries {self.count_of_retries} exceeded {self.count}"
-        )
+            else:
+                self.current_attempt = None  # No more fallback attempts left
+                return None
 
     def __str__(self) -> str:
         return f"Current attempt {self.current_attempt} from {len(self.ai_models_behaviour.attempts)}"

lamoom/ai_models/claude/claude_model.py

@@ -1,7 +1,7 @@
 from lamoom.ai_models.ai_model import AI_MODELS_PROVIDER, AIModel
 import logging
 
-from lamoom.ai_models.constants import C_200K
+from lamoom.ai_models.constants import C_200K, C_4K
 from lamoom.responses import AIResponse
 from decimal import Decimal
 from enum import Enum
@@ -27,36 +27,10 @@ class FamilyModel(Enum):
     opus = "Claude 3 Opus"
 
 
-DEFAULT_PRICING = {
-    "price_per_prompt_1k_tokens": Decimal(0.003),
-    "price_per_sample_1k_tokens": Decimal(0.015),
-}
-
-CLAUDE_AI_PRICING = {
-    FamilyModel.haiku.value: {
-        C_200K: {
-            "price_per_prompt_1k_tokens": Decimal(0.00025),
-            "price_per_sample_1k_tokens": Decimal(0.00125),
-        }
-    },
-    FamilyModel.sonnet.value: {
-        C_200K: {
-            "price_per_prompt_1k_tokens": Decimal(0.003),
-            "price_per_sample_1k_tokens": Decimal(0.015),
-        }
-    },
-    FamilyModel.opus.value: {
-        C_200K: {
-            "price_per_prompt_1k_tokens": Decimal(0.015),
-            "price_per_sample_1k_tokens": Decimal(0.075),
-        }
-    },
-}
-
-
 @dataclass(kw_only=True)
 class ClaudeAIModel(AIModel):
     model: str
+    max_tokens: int = C_4K
     api_key: str = None
     provider: AI_MODELS_PROVIDER = AI_MODELS_PROVIDER.CLAUDE
     family: str = None
@@ -92,13 +66,10 @@ def uny_all_messages_with_same_role(self, messages: t.List[dict]) -> t.List[dict
                 result[-1]["content"] += message.get("content")
         return result
 
-    def call(
-        self,
-        messages: t.List[dict],
-        max_tokens: int,
-        client_secrets: dict = {},
-        **kwargs,
-    ) -> AIResponse:
+
+    def call(self, messages: t.List[dict], max_tokens: int, client_secrets: dict = {}, **kwargs) -> AIResponse:
+        max_tokens = min(max_tokens, self.max_tokens)
+        
         common_args = get_common_args(max_tokens)
         kwargs = {
             **common_args,
@@ -152,25 +123,10 @@ def call(
             logger.exception("[CLAUDEAI] failed to handle chat stream", exc_info=e)
             raise RetryableCustomError(f"Claude AI call failed!")
 
+    @property
     def name(self) -> str:
         return self.model
 
-    @property
-    def price_per_prompt_1k_tokens(self) -> Decimal:
-        keys = list(CLAUDE_AI_PRICING[self.family].keys())
-        def_pricing = CLAUDE_AI_PRICING[self.family].get(keys[0])
-        return CLAUDE_AI_PRICING[self.family].get(self.max_tokens, def_pricing)[
-            "price_per_prompt_1k_tokens"
-        ]
-
-    @property
-    def price_per_sample_1k_tokens(self) -> Decimal:
-        keys = list(CLAUDE_AI_PRICING[self.family].keys())
-        def_pricing = CLAUDE_AI_PRICING[self.family].get(keys[0])
-        return CLAUDE_AI_PRICING[self.family].get(self.max_tokens, def_pricing)[
-            "price_per_sample_1k_tokens"
-        ]
-
     def get_params(self) -> t.Dict[str, t.Any]:
         return {
             "model": self.model,

lamoom/ai_models/constants.py

@@ -5,4 +5,4 @@
 
 C_128K = 128_000
 C_200K = 200_000
-C_1M = 1_000_000
+C_1M = 1_000_000