refactor(agents-api): Remove unnecessary stuff

Signed-off-by: Diwank Tomer <[email protected]>
julep-ai · Aug 14, 2024 · c5f3cec · c5f3cec
1 parent 5007d7b
commit c5f3cec
Show file tree

Hide file tree

Showing 8 changed files with 27 additions and 407 deletions.
diff --git a/agents-api/agents_api/activities/dialog_insights.py b/agents-api/agents_api/activities/dialog_insights.py
diff --git a/agents-api/agents_api/activities/relationship_summary.py b/agents-api/agents_api/activities/relationship_summary.py
diff --git a/agents-api/agents_api/activities/salient_questions.py b/agents-api/agents_api/activities/salient_questions.py
diff --git a/agents-api/agents_api/activities/task_steps/__init__.py b/agents-api/agents_api/activities/task_steps/__init__.py
@@ -14,7 +14,9 @@
     UpdateExecutionRequest,
     YieldStep,
 )
-from ...clients.litellm import acompletion
+from ...clients import (
+    litellm,  # We dont directly import `acompletion` so we can mock it
+)
 from ...clients.worker.types import ChatML
 from ...common.protocol.tasks import (
     StepContext,
@@ -57,7 +59,7 @@ async def prompt_step(context: StepContext) -> dict:
 
     settings: dict = context.definition.settings.model_dump()
     # Get settings and run llm
-    response = await acompletion(
+    response = await litellm.acompletion(
         messages=messages,
         **settings,
     )