mirror of
https://github.com/microsoft/autogen.git
synced 2025-07-16 05:21:05 +00:00
b2800d729b
<!-- 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? This PR updates AGS to use the declarative config serialization native to AgentChat. The effect? You can build your teams/artifacts directly in python, run `team.dump_component()` and immediately run it in AGS. Some change details: - Removes ComponentFactory. Instead TeamManager just loads team specs directly using `Team.load_component`. - Some fixes to the UI to simplify drag and drop experience. - Improve layout of nodes... <!-- Please give a short summary of the change and the problem this solves. --> ## Related issue number <!-- For example: "Closes #1234" --> Closes #4439 Closes #5172 ## Checks - [ ] I've included any doc changes needed for https://microsoft.github.io/autogen/. See https://microsoft.github.io/autogen/docs/Contribute#documentation 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. cc @EItanya @nour-bouzid
131 lines
4.3 KiB
Python
131 lines
4.3 KiB
Python
import os
|
|
import asyncio
|
|
import pytest
|
|
from sqlmodel import Session, text, select
|
|
from typing import Generator
|
|
|
|
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
|
|
|
|
|
|
@pytest.fixture
|
|
def test_db() -> Generator[DatabaseManager, None, None]:
|
|
"""Fixture for test database"""
|
|
db_path = "test.db"
|
|
db = DatabaseManager(f"sqlite:///{db_path}")
|
|
db.reset_db()
|
|
# Initialize database instead of create_db_and_tables
|
|
db.initialize_database(auto_upgrade=False)
|
|
yield db
|
|
# 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}")
|
|
|
|
|
|
@pytest.fixture
|
|
def test_user() -> str:
|
|
return "test_user@example.com"
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_team(test_user: str) -> Team:
|
|
"""Create a sample team with proper config"""
|
|
agent = AssistantAgent(
|
|
name="weather_agent",
|
|
model_client=OpenAIChatCompletionClient(
|
|
model="gpt-4",
|
|
),
|
|
)
|
|
|
|
agent_team = RoundRobinGroupChat([agent], termination_condition=TextMentionTermination("TERMINATE"))
|
|
team_component = agent_team.dump_component()
|
|
|
|
return Team(
|
|
user_id=test_user,
|
|
component=team_component.model_dump(),
|
|
)
|
|
|
|
|
|
class TestDatabaseOperations:
|
|
def test_basic_setup(self, test_db: DatabaseManager):
|
|
"""Test basic database setup and connection"""
|
|
with Session(test_db.engine) as session:
|
|
result = session.exec(text("SELECT 1")).first()
|
|
assert result[0] == 1
|
|
result = session.exec(select(1)).first()
|
|
assert result == 1
|
|
|
|
def test_basic_entity_creation(self, test_db: DatabaseManager, sample_team: Team):
|
|
"""Test creating all entity types with proper configs"""
|
|
# Use upsert instead of raw session
|
|
response = test_db.upsert(sample_team)
|
|
assert response.status is True
|
|
|
|
with Session(test_db.engine) as session:
|
|
saved_team = session.get(Team, sample_team.id)
|
|
assert saved_team is not None
|
|
|
|
def test_upsert_operations(self, test_db: DatabaseManager, sample_team: Team):
|
|
"""Test upsert for both create and update scenarios"""
|
|
# Test Create
|
|
response = test_db.upsert(sample_team)
|
|
assert response.status is True
|
|
assert "Created Successfully" in response.message
|
|
|
|
# Test Update
|
|
team_id = sample_team.id
|
|
sample_team.version = "0.0.2"
|
|
response = test_db.upsert(sample_team)
|
|
assert response.status is True
|
|
|
|
# Verify Update
|
|
result = test_db.get(Team, {"id": team_id})
|
|
assert result.status is True
|
|
assert result.data[0].version == "0.0.2"
|
|
|
|
def test_delete_operations(self, test_db: DatabaseManager, sample_team: Team):
|
|
"""Test delete with various filters"""
|
|
# First insert the model
|
|
response = test_db.upsert(sample_team)
|
|
assert response.status is True # Verify insert worked
|
|
|
|
# Get the ID that was actually saved
|
|
team_id = sample_team.id
|
|
|
|
# Test deletion by id
|
|
response = test_db.delete(Team, {"id": team_id})
|
|
assert response.status is True
|
|
assert "Deleted Successfully" in response.message
|
|
|
|
# Verify deletion
|
|
result = test_db.get(Team, {"id": team_id})
|
|
assert len(result.data) == 0
|
|
|
|
def test_initialize_database_scenarios(self):
|
|
"""Test different initialize_database parameters"""
|
|
db_path = "test_init.db"
|
|
db = DatabaseManager(f"sqlite:///{db_path}")
|
|
|
|
try:
|
|
# Test basic initialization
|
|
response = db.initialize_database()
|
|
assert response.status is True
|
|
|
|
# Test with auto_upgrade
|
|
response = db.initialize_database(auto_upgrade=True)
|
|
assert response.status is True
|
|
|
|
finally:
|
|
asyncio.run(db.close())
|
|
db.reset_db()
|
|
if os.path.exists(db_path):
|
|
os.remove(db_path) |