Skip to content

Commit

Permalink
fix: added auto-generated response_message_handler, since poetry does…
Browse files Browse the repository at this point in the history
…n't support custom build steps
  • Loading branch information
adubovik committed Sep 5, 2024
1 parent 4b7caf9 commit 54e3e6f
Show file tree
Hide file tree
Showing 2 changed files with 112 additions and 2 deletions.
3 changes: 1 addition & 2 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -12,5 +12,4 @@ dist
core-data
core-logs
config.json
~*
response_message_handler.py
~*
111 changes: 111 additions & 0 deletions aidial_interceptors_sdk/chat_completion/response_message_handler.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@
"""
CODE-GENERATED from aidial_interceptors_sdk/chat_completion/request_message_handler.py module.
DO NOT MODIFY THIS FILE.
Callbacks to handle messages in the chat completion response.
"""

from abc import ABC
from typing import List

from aidial_sdk.pydantic_v1 import BaseModel

from aidial_interceptors_sdk.chat_completion.element_path import ElementPath
from aidial_interceptors_sdk.chat_completion.helpers import (
traverse_dict_value,
traverse_list,
)
from aidial_interceptors_sdk.utils.not_given import NotGiven


class ResponseMessageHandler(ABC, BaseModel):
class Config:
arbitrary_types_allowed = True

async def on_response_stage(
self, path: ElementPath, stage: dict
) -> List[dict] | dict:
return stage

async def on_response_stages(
self, path: ElementPath, stages: List[dict] | NotGiven | None
) -> List[dict] | NotGiven | None:
return stages

async def on_response_attachment(
self, path: ElementPath, attachment: dict
) -> List[dict] | dict:
"""
Applied to individual attachments: every message and stage attachment
"""

return attachment

async def on_response_attachments(
self, path: ElementPath, attachments: List[dict] | NotGiven | None
) -> List[dict] | NotGiven | None:
"""
Applied to list of attachments: message attachments and stage attachments
"""
return attachments

async def on_response_state(
self, path: ElementPath, state: dict | NotGiven | None
) -> dict | NotGiven | None:
return state

async def on_response_custom_content(
self, path: ElementPath, custom_content: dict | NotGiven | None
) -> dict | NotGiven | None:
return custom_content

async def traverse_response_message(
self, path: ElementPath, message: dict
) -> dict:

async def apply_on_attachments(
path: ElementPath, attachments: List[dict] | NotGiven | None
) -> List[dict] | NotGiven | None:
attachments = await traverse_list(
path.with_attachment_idx,
attachments,
self.on_response_attachment,
)
return await self.on_response_attachments(path, attachments)

async def apply_on_stage(
path: ElementPath, stage: dict
) -> List[dict] | dict:
stage = await traverse_dict_value(
path, stage, "attachments", apply_on_attachments
)

if path.stage_idx is not None and path.choice_ctx is not None:
mapper = path.choice_ctx.stage_index_mapper
stage["index"] = mapper(path.stage_idx)

return await self.on_response_stage(path, stage)

async def apply_on_stages(
path: ElementPath, stages: List[dict] | NotGiven | None
) -> List[dict] | NotGiven | None:
stages = await traverse_list(
path.with_stage_idx, stages, apply_on_stage
)
return await self.on_response_stages(path, stages)

async def apply_on_custom_content(
path: ElementPath, cc: dict | NotGiven | None
) -> dict | NotGiven | None:
cc = await traverse_dict_value(
path, cc, "state", self.on_response_state
)
cc = await traverse_dict_value(
path, cc, "attachments", apply_on_attachments
)
cc = await traverse_dict_value(path, cc, "stages", apply_on_stages)
return await self.on_response_custom_content(path, cc)

return await traverse_dict_value(
path, message, "custom_content", apply_on_custom_content
)

0 comments on commit 54e3e6f

Please sign in to comment.