mirror of https://github.com/microsoft/autogen.git
Fix AGS Cascading Delete Issue (#5804)
<!-- Thank you for your contribution! Please review https://microsoft.github.io/autogen/docs/Contribute before opening a pull request. --> <!-- Please add a reviewer to the assignee section when you create a PR. If you don't have the access to it, we will shortly find a reviewer and assign them to your PR. --> ## Why are these changes needed? > Hey Victor, this is maybe a bug, but when a session is delete, runs and messages for that session are not deleted, any reason why to keep them? @husseinmozannar The main fix is to add a pragma that ensures SQL lite enforces foreign key constraints. Also needed to update error messages for autoupgrade of databases. Also adds a test for cascade deletes and for parts of teammanager With this fix, - Messages get deleted when the run is deleted - Runs get deleted when sessiosn are deleted - Sessions get deleted when a team is deleted <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number <!-- For example: "Closes #1234" --> ## Checks - [ ] I've included any doc changes needed for <https://microsoft.github.io/autogen/>. See <https://github.com/microsoft/autogen/blob/main/CONTRIBUTING.md> to build and test documentation locally. - [ ] I've added tests (if relevant) corresponding to the changes introduced in this PR. - [ ] I've made sure all auto checks have passed.
This commit is contained in:
Victor Dibia
GitHub
parent
8c9961ecba
commit
679a9357f8
|
|
@ -51,6 +51,10 @@ class DatabaseManager:
|
|||
return Response(message="Database initialization already in progress", status=False)
|
||||
|
||||
try:
|
||||
# Enable foreign key constraints for SQLite
|
||||
if "sqlite" in str(self.engine.url):
|
||||
with self.engine.connect() as conn:
|
||||
conn.execute(text("PRAGMA foreign_keys=ON"))
|
||||
inspector = inspect(self.engine)
|
||||
tables_exist = inspector.get_table_names()
|
||||
if not tables_exist:
|
||||
|
|
@ -221,6 +225,8 @@ class DatabaseManager:
|
|||
|
||||
with Session(self.engine) as session:
|
||||
try:
|
||||
if "sqlite" in str(self.engine.url):
|
||||
session.exec(text("PRAGMA foreign_keys=ON"))
|
||||
statement = select(model_class)
|
||||
if filters:
|
||||
conditions = [getattr(model_class, col) == value for col, value in filters.items()]
|
||||
|
|
|
|||
|
|
@ -466,7 +466,10 @@ datefmt = %H:%M:%S
|
|||
if self.upgrade_schema():
|
||||
return True, "Schema was automatically upgraded"
|
||||
else:
|
||||
return False, "Automatic schema upgrade failed"
|
||||
return (
|
||||
False,
|
||||
"Automatic schema upgrade failed. You are seeing this message because there were differences in your current database schema and the most recent version of the Autogen Studio app database. You can ignore the error, or specifically, you can install AutoGen Studio in a new path `autogenstudio ui --appdir <new path>`.",
|
||||
)
|
||||
|
||||
return False, status
|
||||
|
||||
|
|
|
|||
|
|
@ -7,7 +7,8 @@ from uuid import UUID, uuid4
|
|||
|
||||
from autogen_core import ComponentModel
|
||||
from pydantic import ConfigDict
|
||||
from sqlalchemy import ForeignKey, Integer
|
||||
from sqlalchemy import UUID as SQLAlchemyUUID
|
||||
from sqlalchemy import ForeignKey, Integer, String
|
||||
from sqlmodel import JSON, Column, DateTime, Field, SQLModel, func
|
||||
|
||||
from .types import GalleryConfig, MessageConfig, MessageMeta, SettingsConfig, TeamResult
|
||||
|
|
@ -46,7 +47,9 @@ class Message(SQLModel, table=True):
|
|||
session_id: Optional[int] = Field(
|
||||
default=None, sa_column=Column(Integer, ForeignKey("session.id", ondelete="CASCADE"))
|
||||
)
|
||||
run_id: Optional[UUID] = Field(default=None, foreign_key="run.id")
|
||||
run_id: Optional[UUID] = Field(
|
||||
default=None, sa_column=Column(SQLAlchemyUUID, ForeignKey("run.id", ondelete="CASCADE"))
|
||||
)
|
||||
|
||||
message_meta: Optional[Union[MessageMeta, dict]] = Field(default={}, sa_column=Column(JSON))
|
||||
|
||||
|
|
@ -81,7 +84,7 @@ class Run(SQLModel, table=True):
|
|||
|
||||
__table_args__ = {"sqlite_autoincrement": True}
|
||||
|
||||
id: UUID = Field(default_factory=uuid4, primary_key=True, index=True)
|
||||
id: UUID = Field(default_factory=uuid4, sa_column=Column(SQLAlchemyUUID, primary_key=True, index=True, unique=True))
|
||||
created_at: datetime = Field(
|
||||
default_factory=datetime.now, sa_column=Column(DateTime(timezone=True), server_default=func.now())
|
||||
)
|
||||
|
|
|
|||
|
|
@ -0,0 +1,198 @@
|
|||
{
|
||||
"provider": "autogen_agentchat.teams.RoundRobinGroupChat",
|
||||
"component_type": "team",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "A team that runs a group chat with participants taking turns in a round-robin fashion\n to publish a message to all.",
|
||||
"label": "RoundRobinGroupChat",
|
||||
"config": {
|
||||
"participants": [
|
||||
{
|
||||
"provider": "autogen_agentchat.agents.AssistantAgent",
|
||||
"component_type": "agent",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An agent that provides assistance with tool use.",
|
||||
"label": "AssistantAgent",
|
||||
"config": {
|
||||
"name": "planner_agent",
|
||||
"model_client": {
|
||||
"provider": "autogen_ext.models.openai.OpenAIChatCompletionClient",
|
||||
"component_type": "model",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "Chat completion client for OpenAI hosted models.",
|
||||
"label": "OpenAIChatCompletionClient",
|
||||
"config": {
|
||||
"model": "gpt-4"
|
||||
}
|
||||
},
|
||||
"tools": [],
|
||||
"handoffs": [],
|
||||
"model_context": {
|
||||
"provider": "autogen_core.model_context.UnboundedChatCompletionContext",
|
||||
"component_type": "chat_completion_context",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An unbounded chat completion context that keeps a view of the all the messages.",
|
||||
"label": "UnboundedChatCompletionContext",
|
||||
"config": {}
|
||||
},
|
||||
"description": "A helpful assistant that can plan trips.",
|
||||
"system_message": "You are a helpful assistant that can suggest a travel plan for a user based on their request. Respond with a single sentence",
|
||||
"model_client_stream": false,
|
||||
"reflect_on_tool_use": false,
|
||||
"tool_call_summary_format": "{result}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"provider": "autogen_agentchat.agents.AssistantAgent",
|
||||
"component_type": "agent",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An agent that provides assistance with tool use.",
|
||||
"label": "AssistantAgent",
|
||||
"config": {
|
||||
"name": "local_agent",
|
||||
"model_client": {
|
||||
"provider": "autogen_ext.models.openai.OpenAIChatCompletionClient",
|
||||
"component_type": "model",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "Chat completion client for OpenAI hosted models.",
|
||||
"label": "OpenAIChatCompletionClient",
|
||||
"config": {
|
||||
"model": "gpt-4"
|
||||
}
|
||||
},
|
||||
"tools": [],
|
||||
"handoffs": [],
|
||||
"model_context": {
|
||||
"provider": "autogen_core.model_context.UnboundedChatCompletionContext",
|
||||
"component_type": "chat_completion_context",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An unbounded chat completion context that keeps a view of the all the messages.",
|
||||
"label": "UnboundedChatCompletionContext",
|
||||
"config": {}
|
||||
},
|
||||
"description": "A local assistant that can suggest local activities or places to visit.",
|
||||
"system_message": "You are a helpful assistant that can suggest authentic and interesting local activities or places to visit for a user and can utilize any context information provided. Respond with a single sentence",
|
||||
"model_client_stream": false,
|
||||
"reflect_on_tool_use": false,
|
||||
"tool_call_summary_format": "{result}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"provider": "autogen_agentchat.agents.AssistantAgent",
|
||||
"component_type": "agent",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An agent that provides assistance with tool use.",
|
||||
"label": "AssistantAgent",
|
||||
"config": {
|
||||
"name": "language_agent",
|
||||
"model_client": {
|
||||
"provider": "autogen_ext.models.openai.OpenAIChatCompletionClient",
|
||||
"component_type": "model",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "Chat completion client for OpenAI hosted models.",
|
||||
"label": "OpenAIChatCompletionClient",
|
||||
"config": {
|
||||
"model": "gpt-4"
|
||||
}
|
||||
},
|
||||
"tools": [],
|
||||
"handoffs": [],
|
||||
"model_context": {
|
||||
"provider": "autogen_core.model_context.UnboundedChatCompletionContext",
|
||||
"component_type": "chat_completion_context",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An unbounded chat completion context that keeps a view of the all the messages.",
|
||||
"label": "UnboundedChatCompletionContext",
|
||||
"config": {}
|
||||
},
|
||||
"description": "A helpful assistant that can provide language tips for a given destination.",
|
||||
"system_message": "You are a helpful assistant that can review travel plans, providing feedback on important/critical tips about how best to address language or communication challenges for the given destination. If the plan already includes language tips, you can mention that the plan is satisfactory, with rationale.Respond with a single sentence",
|
||||
"model_client_stream": false,
|
||||
"reflect_on_tool_use": false,
|
||||
"tool_call_summary_format": "{result}"
|
||||
}
|
||||
},
|
||||
{
|
||||
"provider": "autogen_agentchat.agents.AssistantAgent",
|
||||
"component_type": "agent",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An agent that provides assistance with tool use.",
|
||||
"label": "AssistantAgent",
|
||||
"config": {
|
||||
"name": "travel_summary_agent",
|
||||
"model_client": {
|
||||
"provider": "autogen_ext.models.openai.OpenAIChatCompletionClient",
|
||||
"component_type": "model",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "Chat completion client for OpenAI hosted models.",
|
||||
"label": "OpenAIChatCompletionClient",
|
||||
"config": {
|
||||
"model": "gpt-4"
|
||||
}
|
||||
},
|
||||
"tools": [],
|
||||
"handoffs": [],
|
||||
"model_context": {
|
||||
"provider": "autogen_core.model_context.UnboundedChatCompletionContext",
|
||||
"component_type": "chat_completion_context",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "An unbounded chat completion context that keeps a view of the all the messages.",
|
||||
"label": "UnboundedChatCompletionContext",
|
||||
"config": {}
|
||||
},
|
||||
"description": "A helpful assistant that can summarize the travel plan.",
|
||||
"system_message": "You are a helpful assistant that can take in all of the suggestions and advice from the other agents and provide a detailed tfinal travel plan. You must ensure th b at the final plan is integrated and complete. YOUR FINAL RESPONSE MUST BE THE COMPLETE PLAN. When the plan is complete and all perspectives are integrated, you can respond with TERMINATE.Respond with a single sentence",
|
||||
"model_client_stream": false,
|
||||
"reflect_on_tool_use": false,
|
||||
"tool_call_summary_format": "{result}"
|
||||
}
|
||||
}
|
||||
],
|
||||
"termination_condition": {
|
||||
"provider": "autogen_agentchat.base.OrTerminationCondition",
|
||||
"component_type": "termination",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"label": "OrTerminationCondition",
|
||||
"config": {
|
||||
"conditions": [
|
||||
{
|
||||
"provider": "autogen_agentchat.conditions.TextMentionTermination",
|
||||
"component_type": "termination",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "Terminate the conversation if a specific text is mentioned.",
|
||||
"label": "TextMentionTermination",
|
||||
"config": {
|
||||
"text": "TERMINATE"
|
||||
}
|
||||
},
|
||||
{
|
||||
"provider": "autogen_agentchat.conditions.MaxMessageTermination",
|
||||
"component_type": "termination",
|
||||
"version": 1,
|
||||
"component_version": 1,
|
||||
"description": "Terminate the conversation after a maximum number of messages have been exchanged.",
|
||||
"label": "MaxMessageTermination",
|
||||
"config": {
|
||||
"max_messages": 10,
|
||||
"include_agent_event": false
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
@ -19,11 +19,22 @@
|
|||
"execution_count": 1,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/Users/victordibia/projects/hax/autogen/python/packages/autogen-core/src/autogen_core/_component_config.py:252: UserWarning: \n",
|
||||
"⚠️ SECURITY WARNING ⚠️\n",
|
||||
"Loading a FunctionTool from config will execute code to import the provided global imports and and function code.\n",
|
||||
"Only load configs from TRUSTED sources to prevent arbitrary code execution.\n",
|
||||
" instance = component_class._from_config(validated_config) # type: ignore\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"task_result=TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the weather in New York?', type='TextMessage'), ToolCallRequestEvent(source='writing_agent', models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15), content=[FunctionCall(id='call_jcgtAVlBvTFzVpPxKX88Xsa4', arguments='{\"city\":\"New York\"}', name='get_weather')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='writing_agent', models_usage=None, content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_jcgtAVlBvTFzVpPxKX88Xsa4')], type='ToolCallExecutionEvent'), TextMessage(source='writing_agent', models_usage=None, content='The weather in New York is 73 degrees and Sunny.', type='TextMessage'), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=103, completion_tokens=14), content='The current weather in New York is 73 degrees and sunny.', type='TextMessage')], stop_reason='Maximum number of messages 5 reached, current message count: 5') usage='' duration=5.103050947189331\n"
|
||||
"task_result=TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='What is the weather in New York?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=47), metadata={}, content=\"I'm unable to provide real-time weather information. I recommend checking a reliable weather website or app for the current weather in New York. If you have any other questions or need assistance with a different topic, feel free to ask!\", type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=138, completion_tokens=5), metadata={}, content='TERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\") usage='' duration=1.734266996383667\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
|
|
@ -44,12 +55,12 @@
|
|||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"source='user' models_usage=None content='What is the weather in New York?' type='TextMessage'\n",
|
||||
"source='writing_agent' models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15) content=[FunctionCall(id='call_EwdwWogp5jDKdB7t9WGCNjZW', arguments='{\"city\":\"New York\"}', name='get_weather')] type='ToolCallRequestEvent'\n",
|
||||
"source='writing_agent' models_usage=None content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_EwdwWogp5jDKdB7t9WGCNjZW')] type='ToolCallExecutionEvent'\n",
|
||||
"source='writing_agent' models_usage=None content='The weather in New York is 73 degrees and Sunny.' type='TextMessage'\n",
|
||||
"source='writing_agent' models_usage=RequestUsage(prompt_tokens=103, completion_tokens=14) content='The weather in New York is currently 73 degrees and sunny.' type='TextMessage'\n",
|
||||
"task_result=TaskResult(messages=[TextMessage(source='user', models_usage=None, content='What is the weather in New York?', type='TextMessage'), ToolCallRequestEvent(source='writing_agent', models_usage=RequestUsage(prompt_tokens=65, completion_tokens=15), content=[FunctionCall(id='call_EwdwWogp5jDKdB7t9WGCNjZW', arguments='{\"city\":\"New York\"}', name='get_weather')], type='ToolCallRequestEvent'), ToolCallExecutionEvent(source='writing_agent', models_usage=None, content=[FunctionExecutionResult(content='The weather in New York is 73 degrees and Sunny.', call_id='call_EwdwWogp5jDKdB7t9WGCNjZW')], type='ToolCallExecutionEvent'), TextMessage(source='writing_agent', models_usage=None, content='The weather in New York is 73 degrees and Sunny.', type='TextMessage'), TextMessage(source='writing_agent', models_usage=RequestUsage(prompt_tokens=103, completion_tokens=14), content='The weather in New York is currently 73 degrees and sunny.', type='TextMessage')], stop_reason='Maximum number of messages 5 reached, current message count: 5') usage='' duration=1.284574270248413\n"
|
||||
"source='user' models_usage=None metadata={} content='What is the weather in New York?' type='TextMessage'\n",
|
||||
"source='assistant_agent' models_usage=RequestUsage(prompt_tokens=86, completion_tokens=34) metadata={} content=\"I currently don't have access to real-time weather data. You can check a reliable weather website or application for the most current information about the weather in New York.\" type='TextMessage'\n",
|
||||
"source='llm_call_event' models_usage=None metadata={} content='{\"type\": \"LLMCall\", \"messages\": [{\"content\": \"You are a helpful assistant. Solve tasks carefully. When done, say TERMINATE.\", \"role\": \"system\"}, {\"content\": \"What is the weather in New York?\", \"role\": \"user\", \"name\": \"user\"}], \"response\": {\"id\": \"chatcmpl-B76eHFx8OqWQlbnNaK3W6CJtuznmR\", \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"logprobs\": null, \"message\": {\"content\": \"I currently don\\'t have access to real-time weather data. You can check a reliable weather website or application for the most current information about the weather in New York.\", \"refusal\": null, \"role\": \"assistant\", \"audio\": null, \"function_call\": null, \"tool_calls\": null}}], \"created\": 1741033553, \"model\": \"gpt-4o-mini-2024-07-18\", \"object\": \"chat.completion\", \"service_tier\": \"default\", \"system_fingerprint\": \"fp_06737a9306\", \"usage\": {\"completion_tokens\": 34, \"prompt_tokens\": 86, \"total_tokens\": 120, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 0, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0}}}, \"prompt_tokens\": 86, \"completion_tokens\": 34, \"agent_id\": \"assistant_agent/3db465b4-fb0a-4ffd-80df-f4554605d4c0\"}'\n",
|
||||
"source='assistant_agent' models_usage=RequestUsage(prompt_tokens=125, completion_tokens=5) metadata={} content='TERMINATE' type='TextMessage'\n",
|
||||
"source='llm_call_event' models_usage=None metadata={} content='{\"type\": \"LLMCall\", \"messages\": [{\"content\": \"You are a helpful assistant. Solve tasks carefully. When done, say TERMINATE.\", \"role\": \"system\"}, {\"content\": \"What is the weather in New York?\", \"role\": \"user\", \"name\": \"user\"}, {\"content\": \"I currently don\\'t have access to real-time weather data. You can check a reliable weather website or application for the most current information about the weather in New York.\", \"role\": \"assistant\", \"name\": \"assistant_agent\"}], \"response\": {\"id\": \"chatcmpl-B76eIpGDvd7M2nxdgezZuRzOJTQR1\", \"choices\": [{\"finish_reason\": \"stop\", \"index\": 0, \"logprobs\": null, \"message\": {\"content\": \"TERMINATE\", \"refusal\": null, \"role\": \"assistant\", \"audio\": null, \"function_call\": null, \"tool_calls\": null}}], \"created\": 1741033554, \"model\": \"gpt-4o-mini-2024-07-18\", \"object\": \"chat.completion\", \"service_tier\": \"default\", \"system_fingerprint\": \"fp_06737a9306\", \"usage\": {\"completion_tokens\": 5, \"prompt_tokens\": 125, \"total_tokens\": 130, \"completion_tokens_details\": {\"accepted_prediction_tokens\": 0, \"audio_tokens\": 0, \"reasoning_tokens\": 0, \"rejected_prediction_tokens\": 0}, \"prompt_tokens_details\": {\"audio_tokens\": 0, \"cached_tokens\": 0}}}, \"prompt_tokens\": 125, \"completion_tokens\": 5, \"agent_id\": \"assistant_agent/3db465b4-fb0a-4ffd-80df-f4554605d4c0\"}'\n",
|
||||
"task_result=TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='What is the weather in New York?', type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=86, completion_tokens=34), metadata={}, content=\"I currently don't have access to real-time weather data. You can check a reliable weather website or application for the most current information about the weather in New York.\", type='TextMessage'), TextMessage(source='assistant_agent', models_usage=RequestUsage(prompt_tokens=125, completion_tokens=5), metadata={}, content='TERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\") usage='' duration=1.1263680458068848\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
|
|
@ -68,25 +79,14 @@
|
|||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"execution_count": 3,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"[<autogen_agentchat.agents._assistant_agent.AssistantAgent object at 0x122b128d0>]\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"/Users/victordibia/projects/hax/autogen/python/packages/autogen-core/src/autogen_core/_component_config.py:252: UserWarning: \n",
|
||||
"⚠️ SECURITY WARNING ⚠️\n",
|
||||
"Loading a FunctionTool from config will execute code to import the provided global imports and and function code.\n",
|
||||
"Only load configs from TRUSTED sources to prevent arbitrary code execution.\n",
|
||||
" instance = component_class._from_config(validated_config) # type: ignore\n"
|
||||
"[<autogen_agentchat.agents._assistant_agent.AssistantAgent object at 0x1219d3850>]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
|
|
@ -109,16 +109,28 @@
|
|||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 3,
|
||||
"execution_count": 15,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stderr",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"\u001b[32m2025-03-03 12:36:12.391\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mautogenstudio.database.db_manager\u001b[0m:\u001b[36minitialize_database\u001b[0m:\u001b[36m61\u001b[0m - \u001b[1mCreating database tables...\u001b[0m\n",
|
||||
"\u001b[32m2025-03-03 12:36:12.425\u001b[0m | \u001b[1mINFO \u001b[0m | \u001b[36mautogenstudio.database.schema_manager\u001b[0m:\u001b[36m_initialize_alembic\u001b[0m:\u001b[36m154\u001b[0m - \u001b[1mAlembic initialization complete\u001b[0m\n",
|
||||
"INFO [alembic.runtime.migration] Context impl SQLiteImpl.\n",
|
||||
"INFO [alembic.runtime.migration] Will assume non-transactional DDL.\n",
|
||||
"INFO [alembic.autogenerate.compare] Detected type change from NUMERIC() to UUID() on 'message.run_id'\n",
|
||||
"INFO [alembic.autogenerate.compare] Detected type change from NUMERIC() to UUID() on 'run.id'\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"data": {
|
||||
"text/plain": [
|
||||
"Response(message='Database is ready', status=True, data=None)"
|
||||
"Response(message='Database initialized successfully', status=True, data=None)"
|
||||
]
|
||||
},
|
||||
"execution_count": 3,
|
||||
"execution_count": 15,
|
||||
"metadata": {},
|
||||
"output_type": "execute_result"
|
||||
}
|
||||
|
|
@ -136,149 +148,6 @@
|
|||
"dbmanager.initialize_database()"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"\n",
|
||||
"from sqlmodel import Session, text, select\n",
|
||||
"from autogenstudio.datamodel.types import ModelTypes, TeamTypes, AgentTypes, ToolConfig, ToolTypes, OpenAIModelConfig, RoundRobinTeamConfig, MaxMessageTerminationConfig, AssistantAgentConfig, TerminationTypes\n",
|
||||
"\n",
|
||||
"from autogenstudio.datamodel.db import Model, Team, Agent, Tool,LinkTypes\n",
|
||||
"\n",
|
||||
"user_id = \"guestuser@gmail.com\" \n",
|
||||
"\n",
|
||||
"gpt4_model = Model(user_id=user_id, config= OpenAIModelConfig(model=\"gpt-4o-2024-08-06\", model_type=ModelTypes.OPENAI).model_dump() )\n",
|
||||
"\n",
|
||||
"weather_tool = Tool(user_id=user_id, config=ToolConfig(name=\"get_weather\", description=\"Get the weather for a city\", content=\"async def get_weather(city: str) -> str:\\n return f\\\"The weather in {city} is 73 degrees and Sunny.\\\"\",tool_type=ToolTypes.PYTHON_FUNCTION).model_dump() )\n",
|
||||
"\n",
|
||||
"adding_tool = Tool(user_id=user_id, config=ToolConfig(name=\"add\", description=\"Add two numbers\", content=\"async def add(a: int, b: int) -> int:\\n return a + b\", tool_type=ToolTypes.PYTHON_FUNCTION).model_dump() )\n",
|
||||
"\n",
|
||||
"writing_agent = Agent(user_id=user_id,\n",
|
||||
" config=AssistantAgentConfig(\n",
|
||||
" name=\"writing_agent\",\n",
|
||||
" tools=[weather_tool.config],\n",
|
||||
" agent_type=AgentTypes.ASSISTANT,\n",
|
||||
" model_client=gpt4_model.config\n",
|
||||
" ).model_dump()\n",
|
||||
" )\n",
|
||||
"\n",
|
||||
"team = Team(user_id=user_id, config=RoundRobinTeamConfig(\n",
|
||||
" name=\"weather_team\",\n",
|
||||
" participants=[writing_agent.config],\n",
|
||||
" termination_condition=MaxMessageTerminationConfig(termination_type=TerminationTypes.MAX_MESSAGES, max_messages=5).model_dump(),\n",
|
||||
" team_type=TeamTypes.ROUND_ROBIN\n",
|
||||
" ).model_dump()\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"with Session(dbmanager.engine) as session:\n",
|
||||
" session.add(gpt4_model)\n",
|
||||
" session.add(weather_tool)\n",
|
||||
" session.add(adding_tool)\n",
|
||||
" session.add(writing_agent)\n",
|
||||
" session.add(team)\n",
|
||||
" session.commit()\n",
|
||||
"\n",
|
||||
" dbmanager.link(LinkTypes.AGENT_MODEL, writing_agent.id, gpt4_model.id)\n",
|
||||
" dbmanager.link(LinkTypes.AGENT_TOOL, writing_agent.id, weather_tool.id)\n",
|
||||
" dbmanager.link(LinkTypes.AGENT_TOOL, writing_agent.id, adding_tool.id)\n",
|
||||
" dbmanager.link(LinkTypes.TEAM_AGENT, team.id, writing_agent.id)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"2 teams in database\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"all_teams = dbmanager.get(Team)\n",
|
||||
"print(len(all_teams.data), \"teams in database\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
"source": [
|
||||
"## Configuration Manager\n",
|
||||
"\n",
|
||||
"Helper class to mostly import teams/agents/models/tools etc into a database."
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 8,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from autogenstudio.database import ConfigurationManager\n",
|
||||
"\n",
|
||||
"config_manager = ConfigurationManager(dbmanager)\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 9,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"message='Team Created Successfully' status=True data={'id': 4, 'updated_at': datetime.datetime(2024, 12, 15, 15, 52, 21, 674916), 'version': '0.0.1', 'created_at': datetime.datetime(2024, 12, 15, 15, 52, 21, 674910), 'user_id': 'user_id', 'config': {'version': '1.0.0', 'component_type': 'team', 'name': 'weather_team', 'participants': [{'version': '1.0.0', 'component_type': 'agent', 'name': 'writing_agent', 'agent_type': 'AssistantAgent', 'description': None, 'model_client': {'version': '1.0.0', 'component_type': 'model', 'model': 'gpt-4o-2024-08-06', 'model_type': 'OpenAIChatCompletionClient', 'api_key': None, 'base_url': None}, 'tools': [{'version': '1.0.0', 'component_type': 'tool', 'name': 'get_weather', 'description': 'Get the weather for a city', 'content': 'async def get_weather(city: str) -> str:\\n return f\"The weather in {city} is 73 degrees and Sunny.\"', 'tool_type': 'PythonFunction'}], 'system_message': None}], 'team_type': 'RoundRobinGroupChat', 'termination_condition': {'version': '1.0.0', 'component_type': 'termination', 'termination_type': 'MaxMessageTermination', 'max_messages': 5}, 'max_turns': None}}\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"result = await config_manager.import_component(\"team.json\", user_id=\"user_id\", check_exists=True)\n",
|
||||
"print(result)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 10,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"message='Directory import complete' status=True data=[{'component': 'team', 'status': True, 'message': 'Team Created Successfully', 'id': 5}]\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"result = await config_manager.import_directory(\".\", user_id=\"user_id\", check_exists=False)\n",
|
||||
"print(result)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"5 teams in database\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"all_teams = dbmanager.get(Team)\n",
|
||||
"print(len(all_teams.data), \"teams in database\")"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "markdown",
|
||||
"metadata": {},
|
||||
|
|
@ -288,12 +157,12 @@
|
|||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 12,
|
||||
"execution_count": null,
|
||||
"metadata": {},
|
||||
"outputs": [],
|
||||
"source": [
|
||||
"from autogen_agentchat.agents import AssistantAgent\n",
|
||||
"from autogen_agentchat.conditions import TextMentionTermination\n",
|
||||
"from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination\n",
|
||||
"from autogen_agentchat.teams import RoundRobinGroupChat, SelectorGroupChat\n",
|
||||
"from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
|
||||
"\n",
|
||||
|
|
@ -325,36 +194,75 @@
|
|||
" system_message=\"You are a helpful assistant that can take in all of the suggestions and advice from the other agents and provide a detailed tfinal travel plan. You must ensure th b at the final plan is integrated and complete. YOUR FINAL RESPONSE MUST BE THE COMPLETE PLAN. When the plan is complete and all perspectives are integrated, you can respond with TERMINATE.Respond with a single sentence\",\n",
|
||||
")\n",
|
||||
"\n",
|
||||
"termination = TextMentionTermination(\"TERMINATE\")\n",
|
||||
"termination = TextMentionTermination(\"TERMINATE\") | MaxMessageTermination(10)\n",
|
||||
"group_chat = RoundRobinGroupChat(\n",
|
||||
" [planner_agent, local_agent, language_agent, travel_summary_agent], termination_condition=termination\n",
|
||||
")"
|
||||
")\n",
|
||||
"\n",
|
||||
"\n"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 13,
|
||||
"execution_count": 7,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"source='user' models_usage=None content='Plan a 3 day trip to Nepal.' type='TextMessage'\n",
|
||||
"source='planner_agent' models_usage=RequestUsage(prompt_tokens=45, completion_tokens=53) content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.' type='TextMessage'\n",
|
||||
"source='local_agent' models_usage=RequestUsage(prompt_tokens=115, completion_tokens=53) content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.' type='TextMessage'\n",
|
||||
"source='language_agent' models_usage=RequestUsage(prompt_tokens=199, completion_tokens=42) content=\"For your trip to Nepal, it's crucial to learn some phrases in Nepali since English is not widely spoken outside of major cities and tourist areas; even a simple phrasebook or translation app would be beneficial.\" type='TextMessage'\n",
|
||||
"source='travel_summary_agent' models_usage=RequestUsage(prompt_tokens=265, completion_tokens=298) content=\"Day 1: Begin your journey in Kathmandu, where you can visit the historic Durbar Square, a UNESCO World Heritage site that showcases intricate woodcarving and houses the iconic Kasthamandap Temple. From there, proceed to the sacred Pashupatinath Temple, a significant Hindu pilgrimage site on the banks of the holy Bagmati River.\\n\\nDay 2: Embark on an early morning scenic flight over the Everest range. This one-hour flight provides a breathtaking view of the world's highest peak along with other neighboring peaks. Standard flights depart from Tribhuvan International Airport between 6:30 AM to 7:30 AM depending on the weather. Spend the remainder of the day exploring the local markets in Kathmandu, sampling a variety of Nepalese cuisines and shopping for unique souvenirs.\\n\\nDay 3: Finally, take a short flight or drive to Pokhara, the gateway to the Annapurna region. Embark on a guided hike enjoying the stunning backdrop of the Annapurna ranges and the serene Phewa lake.\\n\\nRemember to bring along a phrasebook or translation app, as English is not widely spoken in Nepal, particularly outside of major cities and tourist hotspots. \\n\\nPack comfortable trekking gear, adequate water, medical and emergency supplies. It's also advisable to check on the weather updates, as conditions can change rapidly, particularly in mountainous areas. Enjoy your Nepal expedition!TERMINATE\" type='TextMessage'\n",
|
||||
"TaskResult(messages=[TextMessage(source='user', models_usage=None, content='Plan a 3 day trip to Nepal.', type='TextMessage'), TextMessage(source='planner_agent', models_usage=RequestUsage(prompt_tokens=45, completion_tokens=53), content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.', type='TextMessage'), TextMessage(source='local_agent', models_usage=RequestUsage(prompt_tokens=115, completion_tokens=53), content='I recommend starting your trip in Kathmandu, where you can explore the historic Durbar Square and Pashupatinath Temple, then take a scenic flight over the Everest range, and finish your journey with a stunning hike in the Annapurna region.', type='TextMessage'), TextMessage(source='language_agent', models_usage=RequestUsage(prompt_tokens=199, completion_tokens=42), content=\"For your trip to Nepal, it's crucial to learn some phrases in Nepali since English is not widely spoken outside of major cities and tourist areas; even a simple phrasebook or translation app would be beneficial.\", type='TextMessage'), TextMessage(source='travel_summary_agent', models_usage=RequestUsage(prompt_tokens=265, completion_tokens=298), content=\"Day 1: Begin your journey in Kathmandu, where you can visit the historic Durbar Square, a UNESCO World Heritage site that showcases intricate woodcarving and houses the iconic Kasthamandap Temple. From there, proceed to the sacred Pashupatinath Temple, a significant Hindu pilgrimage site on the banks of the holy Bagmati River.\\n\\nDay 2: Embark on an early morning scenic flight over the Everest range. This one-hour flight provides a breathtaking view of the world's highest peak along with other neighboring peaks. Standard flights depart from Tribhuvan International Airport between 6:30 AM to 7:30 AM depending on the weather. Spend the remainder of the day exploring the local markets in Kathmandu, sampling a variety of Nepalese cuisines and shopping for unique souvenirs.\\n\\nDay 3: Finally, take a short flight or drive to Pokhara, the gateway to the Annapurna region. Embark on a guided hike enjoying the stunning backdrop of the Annapurna ranges and the serene Phewa lake.\\n\\nRemember to bring along a phrasebook or translation app, as English is not widely spoken in Nepal, particularly outside of major cities and tourist hotspots. \\n\\nPack comfortable trekking gear, adequate water, medical and emergency supplies. It's also advisable to check on the weather updates, as conditions can change rapidly, particularly in mountainous areas. Enjoy your Nepal expedition!TERMINATE\", type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")\n"
|
||||
"source='user' models_usage=None metadata={} content='Plan a 3 day trip to Nepal.' type='TextMessage'\n",
|
||||
"source='planner_agent' models_usage=RequestUsage(prompt_tokens=45, completion_tokens=64) metadata={} content='Sure, on your 3-day trip to Nepal, you could start by visiting historical landmarks and local markets in Kathmandu, hike for a day in the Annapurna region to experience breathtaking mountain views and end your trip by exploring the religious sites and wildlife in Lumbini and Chitwan respectively.' type='TextMessage'\n",
|
||||
"source='local_agent' models_usage=RequestUsage(prompt_tokens=125, completion_tokens=86) metadata={} content=\"Absolutely, on the first day, feel free to explore Kathmandu's rich history and culture, visit temples like Swayambhunath and Pashupatinath. On the second day, take a spectacular mountain flight to view Mount Everest and then head to Pokhara to engage in adventure activities or relax by Phewa Lake. On the third day, explore the biodiversity at Chitwan National Park.\" type='TextMessage'\n",
|
||||
"source='language_agent' models_usage=RequestUsage(prompt_tokens=241, completion_tokens=72) metadata={} content=\"While going through your itinerary, I don't see any specific strategies for dealing with language barriers. For Nepal, whilst Nepali is the most widely spoken language, English is understood by many, especially in urban areas and by those working in the travel and hospitality sectors. Carrying a phrasebook or using a translator app could be helpful in more rural areas.\" type='TextMessage'\n",
|
||||
"source='travel_summary_agent' models_usage=RequestUsage(prompt_tokens=336, completion_tokens=220) metadata={} content=\"Your 3-day trip to Nepal is as follows:\\n\\nDay 1: Begin in Kathmandu, exploring the city's rich history and culture. Visit temples such as Swayambhunath and Pashupatinath. Immerse yourself in local markets to experience the vibrant local life.\\n\\nDay 2: As an early bird, embark on a spectacular mountain flight to witness Mount Everest's grandeur, followed by a trip to Pokhara. In Pokhara, you have the opportunity to partake in thrilling adventure activities or unwind by the serene Phewa Lake.\\n\\nDay 3: Dedicate your last day to exploring the biodiversity at Chitwan National Park, which offers an enchanting wildlife experience.\\n\\nLanguage considerations: While a substantial number of people in urban Nepal understand and communicate in English, particularly those in the tourism sector, having a phrasebook or using a translator app can be beneficial in rural regions.\\n\\nThis comprehensive plan ensures that you will enjoy the diversity of cultural, scenic, and wildlife experiences that Nepal has to offer in a short period. TERMINATE.\" type='TextMessage'\n",
|
||||
"TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Plan a 3 day trip to Nepal.', type='TextMessage'), TextMessage(source='planner_agent', models_usage=RequestUsage(prompt_tokens=45, completion_tokens=64), metadata={}, content='Sure, on your 3-day trip to Nepal, you could start by visiting historical landmarks and local markets in Kathmandu, hike for a day in the Annapurna region to experience breathtaking mountain views and end your trip by exploring the religious sites and wildlife in Lumbini and Chitwan respectively.', type='TextMessage'), TextMessage(source='local_agent', models_usage=RequestUsage(prompt_tokens=125, completion_tokens=86), metadata={}, content=\"Absolutely, on the first day, feel free to explore Kathmandu's rich history and culture, visit temples like Swayambhunath and Pashupatinath. On the second day, take a spectacular mountain flight to view Mount Everest and then head to Pokhara to engage in adventure activities or relax by Phewa Lake. On the third day, explore the biodiversity at Chitwan National Park.\", type='TextMessage'), TextMessage(source='language_agent', models_usage=RequestUsage(prompt_tokens=241, completion_tokens=72), metadata={}, content=\"While going through your itinerary, I don't see any specific strategies for dealing with language barriers. For Nepal, whilst Nepali is the most widely spoken language, English is understood by many, especially in urban areas and by those working in the travel and hospitality sectors. Carrying a phrasebook or using a translator app could be helpful in more rural areas.\", type='TextMessage'), TextMessage(source='travel_summary_agent', models_usage=RequestUsage(prompt_tokens=336, completion_tokens=220), metadata={}, content=\"Your 3-day trip to Nepal is as follows:\\n\\nDay 1: Begin in Kathmandu, exploring the city's rich history and culture. Visit temples such as Swayambhunath and Pashupatinath. Immerse yourself in local markets to experience the vibrant local life.\\n\\nDay 2: As an early bird, embark on a spectacular mountain flight to witness Mount Everest's grandeur, followed by a trip to Pokhara. In Pokhara, you have the opportunity to partake in thrilling adventure activities or unwind by the serene Phewa Lake.\\n\\nDay 3: Dedicate your last day to exploring the biodiversity at Chitwan National Park, which offers an enchanting wildlife experience.\\n\\nLanguage considerations: While a substantial number of people in urban Nepal understand and communicate in English, particularly those in the tourism sector, having a phrasebook or using a translator app can be beneficial in rural regions.\\n\\nThis comprehensive plan ensures that you will enjoy the diversity of cultural, scenic, and wildlife experiences that Nepal has to offer in a short period. TERMINATE.\", type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"\n",
|
||||
"result = group_chat.run_stream(task=\"Plan a 3 day trip to Nepal.\")\n",
|
||||
"async for response in result:\n",
|
||||
" print(response)"
|
||||
]
|
||||
},
|
||||
{
|
||||
"cell_type": "code",
|
||||
"execution_count": 11,
|
||||
"metadata": {},
|
||||
"outputs": [
|
||||
{
|
||||
"name": "stdout",
|
||||
"output_type": "stream",
|
||||
"text": [
|
||||
"source='user' models_usage=None metadata={} content='Plan a 3 day trip to Nepal.' type='TextMessage'\n",
|
||||
"source='planner_agent' models_usage=RequestUsage(prompt_tokens=45, completion_tokens=74) metadata={} content='For a 3-day trip to Nepal, start your journey in Kathmandu, visit renowned UNESCO World Heritage Sites including Pashupatinath Temple and Boudhanath Stupa, day two explore the breathtaking city of Pokhara, and on the final day, hike the serene trails of Nagarkot to witness a captivating sunrise over the Himalayas.' type='TextMessage'\n",
|
||||
"source='local_agent' models_usage=RequestUsage(prompt_tokens=135, completion_tokens=57) metadata={} content='Immerse yourself in the local lifestyle by visiting the bustling markets of Kathmandu, explore the ancient, history-rich city of Bhaktapur, take a day hike to the breathtaking monasteries in the mountains and wrap it up with an unforgettable traditional Nepalese dinner.' type='TextMessage'\n",
|
||||
"source='language_agent' models_usage=RequestUsage(prompt_tokens=222, completion_tokens=34) metadata={} content='Consider packing a helpful language guide or translator app to communicate effectively with locals in Nepal since the primary language there is Nepali and not everyone might be fluent in English.' type='TextMessage'\n",
|
||||
"source='travel_summary_agent' models_usage=RequestUsage(prompt_tokens=279, completion_tokens=207) metadata={} content='Your final travel plan for a 3-day trip to Nepal is as follows:\\n\\nDay 1: Arrive in Kathmandu, immerse yourself in the local lifestyle by visiting the bustling markets, Pashupatinath Temple, and Boudhanath Stupa, all renowned UNESCO World Heritage Sites.\\n\\nDay 2: Explore the breathtaking city of Pokhara and the ancient, history-rich city of Bhaktapur. \\n\\nDay 3: Hike the serene trails of Nagarkot to witness a captivating sunrise over the Himalayas, take a day hike to the breathtaking monasteries in the mountains.\\n\\nThroughout your journey, engage with locals and enjoy a traditional Nepalese dinner. As the primary language in Nepal is Nepali, consider packing a helpful language guide or translator app to communicate effectively.\\n\\nEnsure your plan is flexible to accommodate for any unexpected changes or opportunities that present themselves while exploring the wonders of Nepal, remembering that the essence of travel is to appreciate the journey as much as the destination.\\n\\nTERMINATE' type='TextMessage'\n",
|
||||
"TaskResult(messages=[TextMessage(source='user', models_usage=None, metadata={}, content='Plan a 3 day trip to Nepal.', type='TextMessage'), TextMessage(source='planner_agent', models_usage=RequestUsage(prompt_tokens=45, completion_tokens=74), metadata={}, content='For a 3-day trip to Nepal, start your journey in Kathmandu, visit renowned UNESCO World Heritage Sites including Pashupatinath Temple and Boudhanath Stupa, day two explore the breathtaking city of Pokhara, and on the final day, hike the serene trails of Nagarkot to witness a captivating sunrise over the Himalayas.', type='TextMessage'), TextMessage(source='local_agent', models_usage=RequestUsage(prompt_tokens=135, completion_tokens=57), metadata={}, content='Immerse yourself in the local lifestyle by visiting the bustling markets of Kathmandu, explore the ancient, history-rich city of Bhaktapur, take a day hike to the breathtaking monasteries in the mountains and wrap it up with an unforgettable traditional Nepalese dinner.', type='TextMessage'), TextMessage(source='language_agent', models_usage=RequestUsage(prompt_tokens=222, completion_tokens=34), metadata={}, content='Consider packing a helpful language guide or translator app to communicate effectively with locals in Nepal since the primary language there is Nepali and not everyone might be fluent in English.', type='TextMessage'), TextMessage(source='travel_summary_agent', models_usage=RequestUsage(prompt_tokens=279, completion_tokens=207), metadata={}, content='Your final travel plan for a 3-day trip to Nepal is as follows:\\n\\nDay 1: Arrive in Kathmandu, immerse yourself in the local lifestyle by visiting the bustling markets, Pashupatinath Temple, and Boudhanath Stupa, all renowned UNESCO World Heritage Sites.\\n\\nDay 2: Explore the breathtaking city of Pokhara and the ancient, history-rich city of Bhaktapur. \\n\\nDay 3: Hike the serene trails of Nagarkot to witness a captivating sunrise over the Himalayas, take a day hike to the breathtaking monasteries in the mountains.\\n\\nThroughout your journey, engage with locals and enjoy a traditional Nepalese dinner. As the primary language in Nepal is Nepali, consider packing a helpful language guide or translator app to communicate effectively.\\n\\nEnsure your plan is flexible to accommodate for any unexpected changes or opportunities that present themselves while exploring the wonders of Nepal, remembering that the essence of travel is to appreciate the journey as much as the destination.\\n\\nTERMINATE', type='TextMessage')], stop_reason=\"Text 'TERMINATE' mentioned\")\n"
|
||||
]
|
||||
}
|
||||
],
|
||||
"source": [
|
||||
"import json\n",
|
||||
"\n",
|
||||
"# convert to config \n",
|
||||
"config = group_chat.dump_component().model_dump()\n",
|
||||
"# save as json \n",
|
||||
"\n",
|
||||
"with open(\"travel_team.json\", \"w\") as f:\n",
|
||||
" json.dump(config, f, indent=4)\n",
|
||||
"\n",
|
||||
"# load from json\n",
|
||||
"with open(\"travel_team.json\", \"r\") as f:\n",
|
||||
" config = json.load(f)\n",
|
||||
"\n",
|
||||
"group_chat = RoundRobinGroupChat.load_component(config) \n",
|
||||
"result = group_chat.run_stream(task=\"Plan a 3 day trip to Nepal.\") \n",
|
||||
"async for response in result:\n",
|
||||
" print(response)"
|
||||
]
|
||||
}
|
||||
],
|
||||
"metadata": {
|
||||
|
|
|
|||
|
|
@ -1,22 +1,24 @@
|
|||
import os
|
||||
import asyncio
|
||||
import uuid
|
||||
import pytest
|
||||
from sqlmodel import Session, text, select
|
||||
from typing import Generator
|
||||
from pathlib import Path
|
||||
|
||||
from autogenstudio.database import DatabaseManager
|
||||
from autogen_agentchat.agents import AssistantAgent
|
||||
from autogen_agentchat.teams import RoundRobinGroupChat
|
||||
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
||||
from autogen_agentchat.conditions import TextMentionTermination
|
||||
from autogenstudio.datamodel.db import Team
|
||||
from autogenstudio.datamodel.db import Team, Session as SessionModel, Run, Message, RunStatus, MessageConfig
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_db() -> Generator[DatabaseManager, None, None]:
|
||||
"""Fixture for test database"""
|
||||
db_path = "test.db"
|
||||
db = DatabaseManager(f"sqlite:///{db_path}")
|
||||
def test_db(tmp_path) -> Generator[DatabaseManager, None, None]:
|
||||
"""Fixture for test database using temporary paths"""
|
||||
db_path = tmp_path / "test.db"
|
||||
db = DatabaseManager(f"sqlite:///{db_path}", base_dir=tmp_path)
|
||||
db.reset_db()
|
||||
# Initialize database instead of create_db_and_tables
|
||||
db.initialize_database(auto_upgrade=False)
|
||||
|
|
@ -24,11 +26,7 @@ def test_db() -> Generator[DatabaseManager, None, None]:
|
|||
# Clean up
|
||||
asyncio.run(db.close())
|
||||
db.reset_db()
|
||||
try:
|
||||
if os.path.exists(db_path):
|
||||
os.remove(db_path)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to remove test database file: {e}")
|
||||
# No need to manually remove files - tmp_path is cleaned up automatically
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -109,11 +107,70 @@ class TestDatabaseOperations:
|
|||
# Verify deletion
|
||||
result = test_db.get(Team, {"id": team_id})
|
||||
assert len(result.data) == 0
|
||||
|
||||
def test_cascade_delete(self, test_db: DatabaseManager, test_user: str):
|
||||
"""Test all levels of cascade delete"""
|
||||
# Enable foreign keys for SQLite (crucial for cascade delete)
|
||||
with Session(test_db.engine) as session:
|
||||
session.execute(text("PRAGMA foreign_keys=ON"))
|
||||
session.commit()
|
||||
|
||||
def test_initialize_database_scenarios(self):
|
||||
# Test Run -> Message cascade
|
||||
team1 = Team(user_id=test_user, component={"name": "Team1", "type": "team"})
|
||||
test_db.upsert(team1)
|
||||
session1 = SessionModel(user_id=test_user, team_id=team1.id, name="Session1")
|
||||
test_db.upsert(session1)
|
||||
run1_id = uuid.uuid4()
|
||||
test_db.upsert(Run(
|
||||
id=run1_id,
|
||||
user_id=test_user,
|
||||
session_id=session1.id,
|
||||
status=RunStatus.COMPLETE,
|
||||
task=MessageConfig(content="Task1", source="user").model_dump()
|
||||
))
|
||||
test_db.upsert(Message(
|
||||
user_id=test_user,
|
||||
session_id=session1.id,
|
||||
run_id=run1_id,
|
||||
config=MessageConfig(content="Message1", source="assistant").model_dump()
|
||||
))
|
||||
|
||||
test_db.delete(Run, {"id": run1_id})
|
||||
assert len(test_db.get(Message, {"run_id": run1_id}).data) == 0, "Run->Message cascade failed"
|
||||
|
||||
# Test Session -> Run -> Message cascade
|
||||
session2 = SessionModel(user_id=test_user, team_id=team1.id, name="Session2")
|
||||
test_db.upsert(session2)
|
||||
run2_id = uuid.uuid4()
|
||||
test_db.upsert(Run(
|
||||
id=run2_id,
|
||||
user_id=test_user,
|
||||
session_id=session2.id,
|
||||
status=RunStatus.COMPLETE,
|
||||
task=MessageConfig(content="Task2", source="user").model_dump()
|
||||
))
|
||||
test_db.upsert(Message(
|
||||
user_id=test_user,
|
||||
session_id=session2.id,
|
||||
run_id=run2_id,
|
||||
config=MessageConfig(content="Message2", source="assistant").model_dump()
|
||||
))
|
||||
|
||||
test_db.delete(SessionModel, {"id": session2.id})
|
||||
assert len(test_db.get(Run, {"session_id": session2.id}).data) == 0, "Session->Run cascade failed"
|
||||
assert len(test_db.get(Message, {"run_id": run2_id}).data) == 0, "Session->Run->Message cascade failed"
|
||||
|
||||
# Clean up
|
||||
test_db.delete(Team, {"id": team1.id})
|
||||
|
||||
def test_initialize_database_scenarios(self, tmp_path, monkeypatch):
|
||||
"""Test different initialize_database parameters"""
|
||||
db_path = "test_init.db"
|
||||
db = DatabaseManager(f"sqlite:///{db_path}")
|
||||
db_path = tmp_path / "test_init.db"
|
||||
db = DatabaseManager(f"sqlite:///{db_path}", base_dir=tmp_path)
|
||||
|
||||
# Mock the schema manager's check_schema_status to avoid migration issues
|
||||
monkeypatch.setattr(db.schema_manager, "check_schema_status", lambda: (False, None))
|
||||
monkeypatch.setattr(db.schema_manager, "ensure_schema_up_to_date", lambda: True)
|
||||
|
||||
try:
|
||||
# Test basic initialization
|
||||
|
|
@ -126,6 +183,4 @@ class TestDatabaseOperations:
|
|||
|
||||
finally:
|
||||
asyncio.run(db.close())
|
||||
db.reset_db()
|
||||
if os.path.exists(db_path):
|
||||
os.remove(db_path)
|
||||
db.reset_db()
|
||||
|
|
@ -0,0 +1,149 @@
|
|||
import os
|
||||
import json
|
||||
import pytest
|
||||
import asyncio
|
||||
from pathlib import Path
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
from autogenstudio.teammanager import TeamManager
|
||||
from autogenstudio.datamodel.types import TeamResult, EnvironmentVariable
|
||||
from autogen_core import CancellationToken
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_config():
|
||||
"""Create an actual team and dump its configuration"""
|
||||
from autogen_agentchat.agents import AssistantAgent
|
||||
from autogen_agentchat.teams import RoundRobinGroupChat
|
||||
from autogen_ext.models.openai import OpenAIChatCompletionClient
|
||||
from autogen_agentchat.conditions import TextMentionTermination
|
||||
|
||||
agent = AssistantAgent(
|
||||
name="weather_agent",
|
||||
model_client=OpenAIChatCompletionClient(
|
||||
model="gpt-4o-mini",
|
||||
),
|
||||
)
|
||||
|
||||
agent_team = RoundRobinGroupChat(
|
||||
[agent],
|
||||
termination_condition=TextMentionTermination("TERMINATE")
|
||||
)
|
||||
|
||||
# Dump component and return as dict
|
||||
config = agent_team.dump_component()
|
||||
return config.model_dump()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config_file(sample_config, tmp_path):
|
||||
"""Create a temporary config file"""
|
||||
config_path = tmp_path / "test_config.json"
|
||||
with open(config_path, "w") as f:
|
||||
json.dump(sample_config, f)
|
||||
return config_path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config_dir(sample_config, tmp_path):
|
||||
"""Create a temporary directory with multiple config files"""
|
||||
# Create JSON config
|
||||
json_path = tmp_path / "team1.json"
|
||||
with open(json_path, "w") as f:
|
||||
json.dump(sample_config, f)
|
||||
|
||||
# Create YAML config from the same dict
|
||||
import yaml
|
||||
yaml_path = tmp_path / "team2.yaml"
|
||||
# Create a modified copy to verify we can distinguish between them
|
||||
yaml_config = sample_config.copy()
|
||||
yaml_config["label"] = "YamlTeam" # Change a field to identify this as the YAML version
|
||||
with open(yaml_path, "w") as f:
|
||||
yaml.dump(yaml_config, f)
|
||||
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestTeamManager:
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_from_file(self, config_file, sample_config):
|
||||
"""Test loading configuration from a file"""
|
||||
config = await TeamManager.load_from_file(config_file)
|
||||
assert config == sample_config
|
||||
|
||||
# Test file not found
|
||||
with pytest.raises(FileNotFoundError):
|
||||
await TeamManager.load_from_file("nonexistent_file.json")
|
||||
|
||||
# Test unsupported format
|
||||
wrong_format = config_file.with_suffix(".txt")
|
||||
wrong_format.touch()
|
||||
with pytest.raises(ValueError, match="Unsupported file format"):
|
||||
await TeamManager.load_from_file(wrong_format)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_load_from_directory(self, config_dir):
|
||||
"""Test loading all configurations from a directory"""
|
||||
configs = await TeamManager.load_from_directory(config_dir)
|
||||
assert len(configs) == 2
|
||||
|
||||
# Check if at least one team has expected label
|
||||
team_labels = [config.get("label") for config in configs]
|
||||
assert "RoundRobinGroupChat" in team_labels or "YamlTeam" in team_labels
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_team(self, sample_config):
|
||||
"""Test creating a team from config"""
|
||||
team_manager = TeamManager()
|
||||
|
||||
# Mock Team.load_component
|
||||
with patch("autogen_agentchat.base.Team.load_component") as mock_load:
|
||||
mock_team = MagicMock()
|
||||
mock_load.return_value = mock_team
|
||||
|
||||
team = await team_manager._create_team(sample_config)
|
||||
assert team == mock_team
|
||||
mock_load.assert_called_once_with(sample_config)
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_run_stream(self, sample_config):
|
||||
"""Test streaming team execution results"""
|
||||
team_manager = TeamManager()
|
||||
|
||||
# Mock _create_team and team.run_stream
|
||||
with patch.object(team_manager, "_create_team") as mock_create:
|
||||
mock_team = MagicMock()
|
||||
|
||||
# Create some mock messages to stream
|
||||
mock_messages = [MagicMock(), MagicMock()]
|
||||
mock_result = MagicMock() # TaskResult from run
|
||||
mock_messages.append(mock_result) # Last message is the result
|
||||
|
||||
# Set up the async generator for run_stream
|
||||
async def mock_run_stream(*args, **kwargs):
|
||||
for msg in mock_messages:
|
||||
yield msg
|
||||
|
||||
mock_team.run_stream = mock_run_stream
|
||||
mock_create.return_value = mock_team
|
||||
|
||||
# Call run_stream and collect results
|
||||
streamed_messages = []
|
||||
async for message in team_manager.run_stream(
|
||||
task="Test task",
|
||||
team_config=sample_config
|
||||
):
|
||||
streamed_messages.append(message)
|
||||
|
||||
# Verify the team was created
|
||||
mock_create.assert_called_once()
|
||||
|
||||
# Check that we got the expected number of messages +1 for the TeamResult
|
||||
assert len(streamed_messages) == len(mock_messages)
|
||||
|
||||
# Verify the last message is a TeamResult
|
||||
assert isinstance(streamed_messages[-1], type(mock_messages[-1]))
|
||||
|
||||
Loading…
Reference in New Issue