src/agent/base.py

import importlib
import json

from pydantic import BaseModel, Field, ConfigDict, model_validator
from timezonefinder import TimezoneFinder
import pytz

from agno.agent import Agent, RunEvent
from agno.team import Team

from src.infra.logger import get_logger
logger = get_logger(__name__)

check_list = ["description", "instructions", "expected_output", "role", "markdown"]


class Location(BaseModel):
    """Geographical location with latitude and longitude"""
    lat: float = Field(..., description="Latitude of the location", ge=-90, le=90)
    lng: float = Field(..., description="Longitude of the location", ge=-180, le=180)

    model_config = ConfigDict(frozen=True, extra='forbid')


class UserState(BaseModel):
    """User state containing tasks and preferences"""
    user_id: str = Field(None, description="Unique identifier for the user")

    location: Location = Field(..., description="Current location of the user")

    utc_offset: str = Field(
        default="",  # 临时占位符
        description="User's timezone offset (e.g., 'UTC', )"
    )

    @model_validator(mode='after')
    def set_timezone_from_location(self) -> 'UserState':
        """Automatically set timezone based on location coordinates"""
        if self.utc_offset is None and self.location:
            tf = TimezoneFinder()
            tz_name = tf.timezone_at(lat=self.location.lat, lng=self.location.lng)
            timezone = tz_name if tz_name else 'UTC'
            self.utc_offset = pytz.timezone(timezone)
        return self

def get_context(use_state: UserState):
    context = {"lat": use_state.location.lat,
               "lng": use_state.location.lng}
    return f"<current_location> {context} </current_location>"

def get_setting(name):
    name = name.lower()
    try:
        _module = importlib.import_module(f".setting.{name}", package=__package__)
        return {key: getattr(_module, f"{key}", None) for key in check_list}
    except ModuleNotFoundError:
        logger.warning(f"setting module '.setting.{name}' not found.")
        return {}

def creat_agent(name, model, **kwargs):
    prompt_dict = get_setting(name)
    prompt_dict.update(kwargs)
    return Agent(
        name=name,
        model=model,
        **prompt_dict
    )

def creat_team(name, model, members, **kwargs):
    prompt_dict = get_setting(name)
    prompt_dict.update(kwargs)
    return Team(
        name=name,
        model=model,
        members=members,
        **prompt_dict
    )