Update - ingest variant_name, and revision to config when running llm…

… on run time
Agenta-AI · Apr 15, 2024 · 8f41136 · 8f41136
1 parent 338312a
commit 8f41136
Show file tree

Hide file tree

Showing 7 changed files with 12 additions and 9 deletions.
diff --git a/agenta-backend/agenta_backend/models/api/api_models.py b/agenta-backend/agenta_backend/models/api/api_models.py
@@ -35,7 +35,6 @@ class Result(BaseModel):
 
 
 class GetConfigResponse(BaseModel):
-    config_id: Optional[str]
     config_name: str
     current_version: int
     parameters: Dict[str, Any]

diff --git a/agenta-backend/agenta_backend/routers/configs_router.py b/agenta-backend/agenta_backend/routers/configs_router.py
@@ -130,6 +130,7 @@ async def get_config(
                     status_code=400,
                     detail=f"Environment {environment_name} does not deploy base {base_id}",
                 )
+            variant_revision = found_variant_revision.revision
             config = found_variant_revision.config
         elif config_name:
             variants_db = await db_manager.list_variants_for_base(base_db)
@@ -143,16 +144,12 @@ async def get_config(
                     status_code=400,
                     detail=f"Config name {config_name} not found for base {base_id}",
                 )
+            variant_revision = found_variant.revision
             config = found_variant.config
         logger.debug(config.parameters)
         return GetConfigResponse(
-            config_id=str(
-                0
-            ),  # TODO: Remove from the model and regenerate the SDK client
             config_name=config.config_name,
-            current_version=str(
-                0
-            ),  # TODO: remove from teh model and regenerate the SDK client
+            current_version=variant_revision,
             parameters=config.parameters,
         )
     except HTTPException as e:

diff --git a/agenta-backend/agenta_backend/services/app_manager.py b/agenta-backend/agenta_backend/services/app_manager.py
@@ -94,6 +94,7 @@ async def start_variant(
         env_vars = {} if env_vars is None else env_vars
         env_vars.update(
             {
+                "AGENTA_VARIANT_NAME": db_app_variant.variant_name,
                 "AGENTA_VARIANT_ID": str(db_app_variant.id),
                 "AGENTA_BASE_ID": str(db_app_variant.base.id),
                 "AGENTA_APP_ID": str(db_app_variant.app.id),

diff --git a/agenta-cli/agenta/client/backend/types/create_span.py b/agenta-cli/agenta/client/backend/types/create_span.py
@@ -18,6 +18,7 @@ class CreateSpan(pydantic.BaseModel):
     id: str
     app_id: typing.Optional[str]
     variant_id: typing.Optional[str]
+    variant_name: typing.Optional[str]
     inputs: typing.Optional[typing.Dict[str, typing.Any]]
     outputs: typing.Optional[typing.List[str]]
     config: typing.Optional[typing.Dict[str, typing.Any]]

diff --git a/agenta-cli/agenta/client/backend/types/get_config_response.py b/agenta-cli/agenta/client/backend/types/get_config_response.py
@@ -12,7 +12,6 @@
 
 
 class GetConfigResponse(pydantic.BaseModel):
-    config_id: typing.Optional[str]
     config_name: str
     current_version: int
     parameters: typing.Dict[str, typing.Any]

diff --git a/agenta-cli/agenta/sdk/agenta_init.py b/agenta-cli/agenta/sdk/agenta_init.py
@@ -90,6 +90,7 @@ def init(
         self.host = host
         self.app_id = os.environ.get("AGENTA_APP_ID") if app_id is None else app_id
         self.variant_id = os.environ.get("AGENTA_VARIANT_ID")
+        self.variant_name = os.environ.get("AGENTA_VARIANT_NAME")
         self.api_key = api_key
         self.config = Config(base_id=base_id, host=host)
 
@@ -193,7 +194,7 @@ def pull(self, config_name: str = "default", environment_name: str = None):
                     + str(ex)
                 )
         try:
-            self.set(**config.parameters)
+            self.set(**{"current_version": config.current_version, **config.parameters})
         except Exception as ex:
             logger.warning("Failed to set the configuration with error: " + str(ex))
 
@@ -237,6 +238,7 @@ def llm_tracing(max_workers: Optional[int] = None) -> Tracing:
         base_url=singleton.host,
         app_id=singleton.app_id,  # type: ignore
         variant_id=singleton.variant_id,  # type: ignore
+        variant_name=singleton.variant_name,
         api_key=singleton.api_key,
         max_workers=max_workers,
     )
diff --git a/agenta-cli/agenta/sdk/tracing/llm_tracing.py b/agenta-cli/agenta/sdk/tracing/llm_tracing.py
@@ -36,6 +36,7 @@ def __init__(
         base_url: str,
         app_id: str,
         variant_id: str,
+        variant_name: Optional[str] = None,
         api_key: Optional[str] = None,
         max_workers: Optional[int] = None,
     ):
@@ -44,6 +45,7 @@ def __init__(
         self.llm_logger = llm_logger
         self.app_id = app_id
         self.variant_id = variant_id
+        self.variant_name = variant_name
         self.tasks_manager = TaskQueue(
             max_workers if max_workers else 4, logger=llm_logger
         )
@@ -101,6 +103,7 @@ def start_parent_span(
             id=span_id,
             app_id=self.app_id,
             variant_id=self.variant_id,
+            variant_name=self.variant_name,
             inputs=inputs,
             name=name,
             config=config,
@@ -131,6 +134,7 @@ def start_span(
             name=name,
             app_id=self.app_id,
             variant_id=self.variant_id,
+            variant_name=self.variant_name,
             config=config,
             environment=self.active_trace.environment,
             parent_span_id=self.parent_span_id,