mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-14 20:40:45 +00:00
70 lines
2.5 KiB
Python
70 lines
2.5 KiB
Python
|
import openai
|
||
|
|
|
||
|
|
from ..agents.base import BaseChatAgent
|
||
|
|
from ..messages import ChatMessage
|
||
|
|
from ..runtimes import SingleThreadedRuntime
|
||
|
|
|
||
|
|
|
||
|
|
class OpenAIAssistantAgent(BaseChatAgent):
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
name: str,
|
||
|
|
description: str,
|
||
|
|
runtime: SingleThreadedRuntime,
|
||
|
|
client: openai.AsyncClient,
|
||
|
|
assistant_id: str,
|
||
|
|
thread_id: str,
|
||
|
|
) -> None:
|
||
|
|
super().__init__(name, description, runtime)
|
||
|
|
self._client = client
|
||
|
|
self._assistant_id = assistant_id
|
||
|
|
self._thread_id = thread_id
|
||
|
|
self._current_session_window_length = 0
|
||
|
|
|
||
|
|
async def on_chat_message(self, message: ChatMessage) -> ChatMessage:
|
||
|
|
print("---------------")
|
||
|
|
print(f"{self.name} received message from {message.sender}: {message.body}")
|
||
|
|
print("---------------")
|
||
|
|
if message.reset:
|
||
|
|
# Reset the current session window.
|
||
|
|
self._current_session_window_length = 0
|
||
|
|
|
||
|
|
# Save the message to the thread.
|
||
|
|
_ = await self._client.beta.threads.messages.create(
|
||
|
|
thread_id=self._thread_id,
|
||
|
|
content=message.body,
|
||
|
|
role="user",
|
||
|
|
metadata={"sender": message.sender},
|
||
|
|
)
|
||
|
|
self._current_session_window_length += 1
|
||
|
|
|
||
|
|
# If the message is a save_message_only message, return early.
|
||
|
|
if message.save_message_only:
|
||
|
|
return ChatMessage(body="OK", sender=self.name)
|
||
|
|
|
||
|
|
# Create a run and wait until it finishes.
|
||
|
|
run = await self._client.beta.threads.runs.create_and_poll(
|
||
|
|
thread_id=self._thread_id,
|
||
|
|
assistant_id=self._assistant_id,
|
||
|
|
truncation_strategy={
|
||
|
|
"type": "last_messages",
|
||
|
|
"last_messages": self._current_session_window_length,
|
||
|
|
},
|
||
|
|
)
|
||
|
|
|
||
|
|
if run.status != "completed":
|
||
|
|
# TODO: handle other statuses.
|
||
|
|
raise ValueError(f"Run did not complete successfully: {run}")
|
||
|
|
|
||
|
|
# Get the last message from the run.
|
||
|
|
response = await self._client.beta.threads.messages.list(self._thread_id, run_id=run.id, order="desc", limit=1)
|
||
|
|
last_message_content = response.data[0].content
|
||
|
|
|
||
|
|
# TODO: handle array of content.
|
||
|
|
text_content = [content for content in last_message_content if content.type == "text"]
|
||
|
|
if not text_content:
|
||
|
|
raise ValueError(f"Expected text content in the last message: {last_message_content}")
|
||
|
|
|
||
|
|
# TODO: handle multiple text content.
|
||
|
|
return ChatMessage(body=text_content[0].text.value, sender=self.name)
|