agenta 0.52.6__py3-none-any.whl → 0.63.2__py3-none-any.whl

agenta/sdk/evaluations/results.py

@@ -0,0 +1,66 @@
+from typing import Optional, Dict, Any
+from uuid import UUID
+
+from agenta.sdk.utils.client import authed_api
+from agenta.sdk.models.evaluations import EvaluationResult
+
+# TODO: ADD TYPES
+
+
+async def acreate(
+    *,
+    run_id: UUID,
+    scenario_id: UUID,
+    step_key: str,
+    # repeat_idx: str,
+    # timestamp: datetime,
+    # interval: float,
+    #
+    testcase_id: Optional[UUID] = None,
+    trace_id: Optional[str] = None,
+    error: Optional[dict] = None,
+    #
+    flags: Optional[Dict[str, Any]] = None,
+    tags: Optional[Dict[str, Any]] = None,
+    meta: Optional[Dict[str, Any]] = None,
+) -> EvaluationResult:
+    payload = dict(
+        results=[
+            dict(
+                flags=flags,
+                tags=tags,
+                meta=meta,
+                #
+                testcase_id=str(testcase_id) if testcase_id else None,
+                trace_id=trace_id,
+                error=error,
+                #
+                # interval=interval,
+                # timestamp=timestamp,
+                # repeat_idx=repeat_idx,
+                step_key=step_key,
+                run_id=str(run_id),
+                scenario_id=str(scenario_id),
+                #
+                status="success",
+            )
+        ]
+    )
+
+    response = authed_api()(
+        method="POST",
+        endpoint=f"/preview/evaluations/results/",
+        json=payload,
+    )
+
+    try:
+        response.raise_for_status()
+    except:
+        print(response.text)
+        raise
+
+    response = response.json()
+
+    result = EvaluationResult(**response["results"][0])
+
+    return result

agenta/sdk/evaluations/runs.py

@@ -0,0 +1,153 @@
+from typing import Optional, Dict, Any
+from uuid import UUID
+
+from agenta.sdk.utils.client import authed_api
+from agenta.sdk.models.evaluations import EvaluationRun, Target
+
+import agenta as ag
+
+# TODO: ADD TYPES
+
+
+async def afetch(
+    *,
+    run_id: UUID,
+) -> Optional[EvaluationRun]:
+    response = authed_api()(
+        method="GET",
+        endpoint=f"/preview/evaluations/runs/{run_id}",
+    )
+
+    try:
+        response.raise_for_status()
+    except:
+        print(response.text)
+        raise
+
+    response = response.json()
+
+    if (not "count" in response) or (response["count"] == 0) or (not "run" in response):
+        return None
+
+    run = EvaluationRun(**response["run"])
+
+    return run
+
+
+async def acreate(
+    *,
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+    #
+    flags: Optional[Dict[str, Any]] = None,
+    tags: Optional[Dict[str, Any]] = None,
+    meta: Optional[Dict[str, Any]] = None,
+    #
+    query_steps: Optional[Target] = None,
+    testset_steps: Optional[Target] = None,
+    application_steps: Optional[Target] = None,
+    evaluator_steps: Optional[Target] = None,
+    #
+    repeats: Optional[int] = None,
+) -> Optional[EvaluationRun]:
+    payload = dict(
+        evaluation=dict(
+            name=name,
+            description=description,
+            #
+            flags=flags,
+            tags=tags,
+            meta=meta,
+            #
+            data=dict(
+                status="running",
+                query_steps=query_steps,
+                testset_steps=testset_steps,
+                application_steps=application_steps,
+                evaluator_steps=evaluator_steps,
+                repeats=repeats,
+            ),
+            #
+            jit={"testsets": True, "evaluators": False},
+        )
+    )
+
+    response = authed_api()(
+        method="POST",
+        endpoint=f"/preview/simple/evaluations/",
+        json=payload,
+    )
+
+    try:
+        response.raise_for_status()
+    except:
+        print(response.text)
+        raise
+
+    response = response.json()
+
+    if (not "evaluation" in response) or (not "id" in response["evaluation"]):
+        return None
+
+    run_id = UUID(response["evaluation"]["id"])
+
+    return await afetch(run_id=run_id)
+
+
+async def aclose(
+    *,
+    run_id: UUID,
+    #
+    status: Optional[str] = "success",
+) -> Optional[EvaluationRun]:
+    response = authed_api()(
+        method="POST",
+        endpoint=f"/preview/evaluations/runs/{run_id}/close/{status}",
+    )
+
+    try:
+        response.raise_for_status()
+    except:
+        print(response.text)
+        raise
+
+    response = response.json()
+
+    if (not "run" in response) or (not "id" in response["run"]):
+        return None
+
+    run_id = UUID(response["run"]["id"])
+
+    return await afetch(run_id=run_id)
+
+
+async def aurl(
+    *,
+    run_id: UUID,
+) -> str:
+    response = authed_api()(
+        method="GET",
+        endpoint=f"/projects",
+        params={"scope": "project"},
+    )
+
+    try:
+        response.raise_for_status()
+    except:
+        print(response.text)
+        raise
+
+    if len(response.json()) != 1:
+        return None
+
+    project_info = response.json()[0]
+
+    workspace_id = project_info.get("workspace_id")
+    project_id = project_info.get("project_id")
+
+    return (
+        f"{ag.DEFAULT_AGENTA_SINGLETON_INSTANCE.host}"
+        f"/w/{workspace_id}"
+        f"/p/{project_id}"
+        f"/evaluations/results/{run_id}"
+    )

agenta/sdk/evaluations/scenarios.py

@@ -0,0 +1,48 @@
+from typing import Optional, Dict, Any
+from uuid import UUID
+
+from agenta.sdk.utils.client import authed_api
+from agenta.sdk.models.evaluations import EvaluationScenario
+
+# TODO: ADD TYPES
+
+
+async def acreate(
+    *,
+    run_id: UUID,
+    #
+    flags: Optional[Dict[str, Any]] = None,
+    tags: Optional[Dict[str, Any]] = None,
+    meta: Optional[Dict[str, Any]] = None,
+) -> EvaluationScenario:
+    payload = dict(
+        scenarios=[
+            dict(
+                flags=flags,
+                tags=tags,
+                meta=meta,
+                #
+                run_id=str(run_id),
+                #
+                status="success",
+            )
+        ]
+    )
+
+    response = authed_api()(
+        method="POST",
+        endpoint=f"/preview/evaluations/scenarios/",
+        json=payload,
+    )
+
+    try:
+        response.raise_for_status()
+    except:
+        print(response.text)
+        raise
+
+    response = response.json()
+
+    scenario = EvaluationScenario(**response["scenarios"][0])
+
+    return scenario

agenta/sdk/litellm/litellm.py

@@ -193,6 +193,9 @@ def litellm_handler():
 
             span.end()
 
+            # Clean up span from dictionary to prevent memory leak
+            del self.span[litellm_call_id]
+
         def log_failure_event(
             self,
             kwargs,
@@ -221,6 +224,9 @@ def litellm_handler():
 
             span.end()
 
+            # Clean up span from dictionary to prevent memory leak
+            del self.span[litellm_call_id]
+
         async def async_log_stream_event(
             self,
             kwargs,
@@ -321,6 +327,9 @@ def litellm_handler():
 
             span.end()
 
+            # Clean up span from dictionary to prevent memory leak
+            del self.span[litellm_call_id]
+
         async def async_log_failure_event(
             self,
             kwargs,
@@ -349,4 +358,7 @@ def litellm_handler():
 
             span.end()
 
+            # Clean up span from dictionary to prevent memory leak
+            del self.span[litellm_call_id]
+
     return LitellmHandler()

agenta/sdk/litellm/mockllm.py

@@ -2,10 +2,12 @@ from typing import Optional, Protocol, Any
 from os import environ
 from contextlib import contextmanager
 
+import litellm
+
 from agenta.sdk.utils.logging import get_module_logger
 
 from agenta.sdk.litellm.mocks import MOCKS
-from agenta.sdk.context.serving import serving_context
+from agenta.sdk.contexts.routing import RoutingContext
 
 AGENTA_LITELLM_MOCK = environ.get("AGENTA_LITELLM_MOCK") or None
 
@@ -65,18 +67,14 @@ def user_aws_credentials_from(ps: dict):
 
 
 class LitellmProtocol(Protocol):
-    async def acompletion(self, *args: Any, **kwargs: Any) -> Any:
-        ...
-
-
-litellm: Optional[LitellmProtocol] = None  # pylint: disable=invalid-name
+    async def acompletion(self, *args: Any, **kwargs: Any) -> Any: ...
 
 
 async def acompletion(*args, **kwargs):
-    mock = AGENTA_LITELLM_MOCK or serving_context.get().mock
+    mock = AGENTA_LITELLM_MOCK or RoutingContext.get().mock
 
     if mock:
-        log.debug("Mocking litellm: %s.", mock)
+        # log.debug("Mocking litellm: %s.", mock)
 
         if mock not in MOCKS:
             mock = "hello"

agenta/sdk/litellm/mocks/__init__.py

@@ -3,7 +3,7 @@ from asyncio import sleep
 
 from pydantic import BaseModel
 
-import agenta as ag
+from agenta.sdk.decorators.tracing import instrument
 
 
 class MockMessageModel(BaseModel):
@@ -18,7 +18,7 @@ class MockResponseModel(BaseModel):
     choices: list[MockChoiceModel]
 
 
-@ag.instrument()
+@instrument()
 def hello_mock_response(*args, **kwargs) -> MockResponseModel:
     return MockResponseModel(
         choices=[
@@ -31,7 +31,7 @@ def hello_mock_response(*args, **kwargs) -> MockResponseModel:
     )
 
 
-@ag.instrument()
+@instrument()
 def chat_mock_response(*args, **kwargs) -> MockResponseModel:
     return MockResponseModel(
         choices=[
@@ -45,7 +45,7 @@ def chat_mock_response(*args, **kwargs) -> MockResponseModel:
     )
 
 
-@ag.instrument()
+@instrument()
 def delay_mock_response(*args, **kwargs) -> MockResponseModel:
     sleep(2)
 
@@ -60,7 +60,7 @@ def delay_mock_response(*args, **kwargs) -> MockResponseModel:
     )
 
 
-@ag.instrument()
+@instrument()
 def capital_mock_response(*args, **kwargs) -> MockResponseModel:
     country = kwargs.get("messages", [{}, {}])[1].get(
         "content", "What is the capital of _____?"

agenta/sdk/managers/applications.py

@@ -0,0 +1,304 @@
+from typing import Dict, Any, Callable, Optional
+from uuid import uuid4, UUID
+
+from agenta.sdk.utils.client import authed_api
+from agenta.sdk.decorators.running import auto_workflow, is_workflow
+from agenta.sdk.models.workflows import (
+    ApplicationRevision,
+    #
+    ApplicationRevisionResponse,
+    #
+    LegacyApplicationFlags,
+    LegacyApplicationData,
+    LegacyApplicationCreate,
+    LegacyApplicationEdit,
+    #
+    LegacyApplicationResponse,
+    #
+    Reference,
+)
+
+from agenta.sdk.utils.references import get_slug_from_name_and_id
+
+
+async def _retrieve_application(
+    application_id: Optional[UUID] = None,
+    application_slug: Optional[str] = None,
+    application_revision_id: Optional[UUID] = None,
+    application_revision_slug: Optional[str] = None,
+) -> Optional[ApplicationRevision]:
+    payload = {
+        "application_ref": (
+            {
+                "id": str(application_id) if application_id else None,
+                "slug": str(application_slug),
+            }
+            if application_id or application_slug
+            else None
+        ),
+        "application_revision_ref": (
+            {
+                "id": (
+                    str(application_revision_id) if application_revision_id else None
+                ),
+                "slug": application_revision_slug,
+            }
+            if application_revision_id or application_revision_slug
+            else None
+        ),
+    }
+
+    # print(" --- payload:", payload)
+
+    response = authed_api()(
+        method="POST",
+        endpoint=f"/preview/legacy/applications/revisions/retrieve",
+        json=payload,
+    )
+    response.raise_for_status()
+
+    application_revision_response = ApplicationRevisionResponse(**response.json())
+
+    application_revision = application_revision_response.application_revision
+
+    # print(" --- application_revision:", application_revision)
+
+    return application_revision
+
+
+async def aretrieve(
+    application_revision_id: Optional[UUID] = None,
+) -> Optional[ApplicationRevision]:
+    # print("\n--------- RETRIEVE APPLICATION")
+
+    response = await _retrieve_application(
+        application_revision_id=application_revision_id,
+    )
+
+    return response
+
+
+async def aupsert(
+    *,
+    application_id: Optional[UUID] = None,
+    application_slug: Optional[str] = None,
+    application_revision_id: Optional[UUID] = None,
+    application_revision_slug: Optional[str] = None,
+    #
+    handler: Callable,
+    script: Optional[str] = None,
+    parameters: Optional[Dict[str, Any]] = None,
+    #
+    name: Optional[str] = None,
+    description: Optional[str] = None,
+) -> Optional[UUID]:
+    # print("\n---------   UPSERT APPLICATION")
+    try:
+        if not is_workflow(handler):
+            application_workflow = auto_workflow(
+                handler,
+                #
+                script=script,
+                parameters=parameters,
+                #
+                name=name,
+                description=description,
+            )
+        else:
+            application_workflow = handler
+
+        req = await application_workflow.inspect()
+
+        legacy_application_flags = LegacyApplicationFlags(**req.flags)
+
+        legacy_application_data = LegacyApplicationData(
+            **(
+                req.interface.model_dump(mode="json", exclude_none=True)
+                if req and req.interface
+                else {}
+            ),
+            **(
+                req.configuration.model_dump(mode="json", exclude_none=True)
+                if req and req.configuration
+                else {}
+            ),
+        )
+
+        # print(
+        #     " ---:", legacy_application_data.model_dump(mode="json", exclude_none=True)
+        # )
+
+        retrieve_response = None
+
+        if req.references is not None:
+            _application_revision_ref = req.references.get("application_revision", {})
+            if isinstance(_application_revision_ref, Reference):
+                _application_revision_ref = _application_revision_ref.model_dump(
+                    mode="json",
+                    exclude_none=True,
+                )
+            if not isinstance(_application_revision_ref, dict):
+                _application_revision_ref = {}
+            _application_revision_id = _application_revision_ref.get("id")
+            _application_revision_slug = _application_revision_ref.get("slug")
+
+            application_revision_id = (
+                application_revision_id or _application_revision_id
+            )
+            application_revision_slug = (
+                application_revision_slug or _application_revision_slug
+            )
+
+            _application_ref = req.references.get("application", {})
+            if isinstance(_application_ref, Reference):
+                _application_ref = _application_ref.model_dump(
+                    mode="json",
+                    exclude_none=True,
+                )
+            if not isinstance(_application_ref, dict):
+                _application_ref = {}
+            _application_id = _application_ref.get("id")
+            _application_slug = _application_ref.get("slug")
+
+            application_id = application_id or _application_id
+            application_slug = application_slug or _application_slug
+
+            revision = req.data.revision if req and req.data else None
+            if revision:
+                name = name or revision.get("name")
+                description = description or revision.get("description")
+
+        name = (
+            name or req.data.revision.get("name")
+            if req and req.data and req.data.revision
+            else None
+        )
+
+        description = (
+            description or req.data.revision.get("description")
+            if req and req.data and req.data.revision
+            else None
+        )
+
+        application_slug = (
+            application_slug
+            or get_slug_from_name_and_id(
+                name=name,
+                id=application_id or uuid4(),
+            )
+            if name
+            else uuid4().hex[-12:]
+        )
+
+        # print(
+        #     application_id,
+        #     application_slug,
+        #     application_revision_id,
+        #     application_revision_slug,
+        # )
+
+        if application_revision_id or application_revision_slug:
+            retrieve_response = await _retrieve_application(
+                application_revision_id=application_revision_id,
+                application_revision_slug=application_revision_slug,
+            )
+        elif application_id or application_slug:
+            retrieve_response = await _retrieve_application(
+                application_id=application_id,
+                application_slug=application_slug,
+            )
+
+    except Exception as e:
+        print("[ERROR]: Failed to prepare application:", e)
+        return None
+
+    # print("Retrieve response:", retrieve_response)
+
+    if retrieve_response and retrieve_response.id and retrieve_response.application_id:
+        application_id = retrieve_response.application_id
+        # print(" --- Updating application...", application_id)
+        application_edit_request = LegacyApplicationEdit(
+            id=application_id,
+            #
+            name=name,
+            description=description,
+            #
+            flags=legacy_application_flags,
+            #
+            data=legacy_application_data,
+        )
+
+        # print(" --- application_edit_request:", application_edit_request)
+
+        response = authed_api()(
+            method="PUT",
+            endpoint=f"/preview/legacy/applications/{application_id}",
+            json={
+                "application": application_edit_request.model_dump(
+                    mode="json",
+                    exclude_none=True,
+                )
+            },
+        )
+
+        # print(" --- response:", response.status_code, response.text)
+
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            print("[ERROR]: Failed to update application:", e)
+            return None
+
+    else:
+        # print(" --- Creating application...")
+        application_create_request = LegacyApplicationCreate(
+            slug=application_slug or uuid4().hex[-12:],
+            #
+            name=name,
+            description=description,
+            #
+            flags=legacy_application_flags,
+            #
+            data=legacy_application_data,
+        )
+
+        # print(" --- application_create_request:", application_create_request)
+
+        response = authed_api()(
+            method="POST",
+            endpoint="/preview/legacy/applications/",
+            json={
+                "application": application_create_request.model_dump(
+                    mode="json",
+                    exclude_none=True,
+                )
+            },
+        )
+
+        # print(" --- response:", response.status_code, response.text)
+
+        try:
+            response.raise_for_status()
+        except Exception as e:
+            print("[ERROR]: Failed to create application:", e)
+            return None
+
+    application_response = LegacyApplicationResponse(**response.json())
+
+    application = application_response.application
+
+    if not application or not application.id:
+        return None
+
+    # print(" --- application:", application)
+
+    application_revision = await _retrieve_application(
+        application_id=application.id,
+    )
+
+    if not application_revision or not application_revision.id:
+        return None
+
+    # print(application_revision, "----------")
+
+    return application_revision.id

agenta/sdk/managers/config.py

@@ -7,7 +7,7 @@ from pydantic import BaseModel
 
 from agenta.sdk.utils.logging import get_module_logger
 from agenta.sdk.managers.shared import SharedManager
-from agenta.sdk.context.serving import serving_context
+from agenta.sdk.contexts.routing import RoutingContext
 
 T = TypeVar("T", bound=BaseModel)
 
@@ -45,7 +45,7 @@ class ConfigManager:
             Only one of these should be provided.
         """
 
-        context = serving_context.get()
+        context = RoutingContext.get()
 
         parameters = context.parameters