add: nebius support (#37)

Author: maisiukartyom

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

@@ -19,6 +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 NEBIUS_API_KEY=${{ secrets.NEBIUS_API_KEY }} >> .env
         cat .env
 
     - name: Install dependencies

lamoom/prompt/lamoom.py

@@ -38,6 +38,7 @@ class Lamoom:
     claude_key: str = None
     gemini_key: str = None
     azure_keys: t.Dict[str, str] = None
+    nebius_key: str = None
     secrets: Secrets = None
 
     clients = {}
@@ -65,6 +66,9 @@ def __post_init__(self):
         if not self.claude_key and self.secrets.CLAUDE_API_KEY:
             logger.debug(f"Using Claude API key from secrets")
             self.claude_key = self.secrets.CLAUDE_API_KEY
+        if not self.nebius_key and self.secrets.NEBIUS_API_KEY:
+            logger.debug(f"Using Nebius API key from secrets")
+            self.nebius_key = self.secrets.NEBIUS_API_KEY
         self.service = LamoomService()
         if self.openai_key:
             self.clients[AI_MODELS_PROVIDER.OPENAI] = {
@@ -85,6 +89,8 @@ def __post_init__(self):
             self.clients[AI_MODELS_PROVIDER.CLAUDE] = {"api_key": self.claude_key}
         if self.gemini_key:
             self.clients[AI_MODELS_PROVIDER.GEMINI] = {"api_key": self.gemini_key}
+        if self.nebius_key:
+            self.clients[AI_MODELS_PROVIDER.NEBIUS] = {"api_key": self.nebius_key}
         self.worker = SaveWorker()
 
     def create_test(
@@ -147,6 +153,10 @@ def extract_provider_name(self, model: str) -> dict:
 
         if "azure" in parts[0].lower() and len(parts) == 3:
             model_provider, realm, model_name = parts
+        elif "nebius" in parts[0].lower() and len(parts) == 3:
+            model_provider = parts[0] 
+            model_name = f"{parts[1]}/{parts[2]}"
+            realm = None
         else:
             model_provider, model_name = parts
             realm = None

lamoom/settings.py

@@ -44,6 +44,7 @@ class Secrets:
     OPENAI_API_KEY: str = field(default_factory=lambda: os.getenv("OPENAI_API_KEY"))
     CLAUDE_API_KEY: str = field(default_factory=lambda: os.getenv("CLAUDE_API_KEY"))
     GEMINI_API_KEY: str = field(default_factory=lambda: os.getenv("GEMINI_API_KEY"))
+    NEBIUS_API_KEY: str = field(default_factory=lambda: os.getenv("NEBIUS_API_KEY"))
     OPENAI_ORG: str = field(default_factory=lambda: os.getenv("OPENAI_ORG"))
     azure_keys: dict = field(
         default_factory=lambda: json.loads(

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "lamoom"
-version = "0.1.34"
+version = "0.1.35"
 description = ""
 authors = ["Lamoom Engineering Team <engineering@lamoom.com>"]
 readme = "README.md"

tests/prompts/test_model.py

@@ -26,6 +26,6 @@ def test_model(client):
     prompt = Prompt(id=prompt_id) 
     prompt.add("{text}", role='user')
 
-    result = client.call(prompt.id, context, "openai/o3-mini")
+    result = client.call(prompt.id, context, "nebius/deepseek-ai/DeepSeek-R1")
 
     assert result.content

tests/test_integrational.py

@@ -54,4 +54,4 @@ def test_loading_prompt_from_service(client: Lamoom):
     result = client.call(prompt.id, context, "azure/useast/gpt-4o")
 
     # should call the prompt with music
-    assert result.prompt.messages[-1] == {'role': 'user', 'content': 'test1\ntest2'} 
+    assert result.prompt.messages[-1] == {'role': 'user', 'content': 'music1\nmusic2'}