Skip to content

Commit

Permalink
fix: Fix imports
Browse files Browse the repository at this point in the history
  • Loading branch information
whiterabbit1983 committed Jul 31, 2024
1 parent 7970748 commit 88e7131
Show file tree
Hide file tree
Showing 3 changed files with 25 additions and 11 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,6 @@
wrap_in_class,
)


valid_transitions = {
# Start state
"init": ["wait", "error", "step", "cancelled"],
Expand Down
25 changes: 18 additions & 7 deletions agents-api/agents_api/routers/sessions/session.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,12 @@
from openai.types.chat.chat_completion import ChatCompletion
from pydantic import UUID4

from ...autogen.openapi_model import DocIds, InputChatMLMessage, Tool
from ...autogen.openapi_model import (
CreateEntryRequest,
DocIds,
InputChatMLMessage,
Tool,
)
from ...clients.embed import embed
from ...clients.temporal import run_summarization_task, run_truncation_task
from ...clients.worker.types import ChatML
Expand All @@ -38,7 +43,6 @@
)
from ...models.entry.create_entries import create_entries
from ...models.session.get_cached_response import get_cached_response
from ...models.session.prepare_chat_context import prepare_chat_context
from ...models.session.prepare_session_data import prepare_session_data
from ...models.session.set_cached_response import set_cached_response
from .exceptions import InputTooBigError
Expand Down Expand Up @@ -208,7 +212,9 @@ async def run(
# TODO: implement locking at some point

# Get session data
session_data = get_session_data(self.developer_id, self.session_id)
session_data = prepare_session_data(
developer_id=self.developer_id, session_id=self.session_id
)
if session_data is None:
raise SessionNotFoundError(self.developer_id, self.session_id)

Expand Down Expand Up @@ -502,15 +508,20 @@ async def backward(
entries: list[Entry] = []
for m in new_input:
entries.append(
Entry(
session_id=self.session_id,
CreateEntryRequest(
role=m.role,
content=m.content,
name=m.name,
)
)

entries.append(new_entry)
entries.append(
CreateEntryRequest(
role=new_entry.role,
content=new_entry.content,
name=new_entry.name,
)
)
bg_task = None

if (
Expand All @@ -524,7 +535,7 @@ async def backward(
else:
raise PromptTooBigError(total_tokens, final_settings.token_budget)

create_entries_query(
create_entries(
developer_id=self.developer_id, session_id=self.session_id, data=entries
)

Expand Down
10 changes: 7 additions & 3 deletions agents-api/agents_api/routers/tasks/routers.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,10 +24,14 @@
from agents_api.clients.temporal import run_task_execution_workflow
from agents_api.common.protocol.tasks import ExecutionInput
from agents_api.dependencies.developer_id import get_developer_id
from agents_api.models.execution.create_execution import create_execution as create_execution_query
from agents_api.models.execution.get_execution import get_execution as get_execution_query
from agents_api.models.execution.create_execution import (
create_execution as create_execution_query,
)
from agents_api.models.execution.get_execution import (
get_execution as get_execution_query,
)
from agents_api.models.execution.get_execution_transition import (
get_execution_transition_query,
get_execution_transition as get_execution_transition_query,
)
from agents_api.models.execution.list_execution_transitions import (
list_execution_transitions as list_execution_transitions_query,
Expand Down

0 comments on commit 88e7131

Please sign in to comment.