pycityagent 2.0.0a43__cp311-cp311-macosx_11_0_arm64.whl

pycityagent/simulation/storage/pg.py

@@ -0,0 +1,184 @@
+import asyncio
+import logging
+from typing import Any
+
+import psycopg
+import psycopg.sql
+import ray
+from psycopg.rows import dict_row
+
+from ...utils.decorators import lock_decorator
+from ...utils.pg_query import PGSQL_DICT, TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+
+logger = logging.getLogger("pg")
+
+
+def create_pg_tables(exp_id: str, dsn: str):
+    for table_type, exec_strs in PGSQL_DICT.items():
+        if not table_type == "experiment":
+            table_name = f"socialcity_{exp_id.replace('-', '_')}_{table_type}"
+        else:
+            table_name = f"socialcity_{table_type}"
+        # # debug str
+        # for _str in [f"DROP TABLE IF EXISTS {table_name}"] + [
+        #     _exec_str.format(table_name=table_name) for _exec_str in exec_strs
+        # ]:
+        #     print(_str)
+        with psycopg.connect(dsn) as conn:
+            with conn.cursor() as cur:
+                if not table_type == "experiment":
+                    # delete table
+                    cur.execute(f"DROP TABLE IF EXISTS {table_name}")  # type:ignore
+                    logger.debug(
+                        f"table:{table_name} sql: DROP TABLE IF EXISTS {table_name}"
+                    )
+                    conn.commit()
+                # create table
+                for _exec_str in exec_strs:
+                    exec_str = _exec_str.format(table_name=table_name)
+                    cur.execute(exec_str)
+                    logger.debug(f"table:{table_name} sql: {exec_str}")
+                conn.commit()
+
+
+@ray.remote
+class PgWriter:
+    def __init__(self, exp_id: str, dsn: str):
+        self.exp_id = exp_id
+        self._dsn = dsn
+        self._lock = asyncio.Lock()
+
+    @lock_decorator
+    async def async_write_dialog(self, rows: list[tuple]):
+        _tuple_types = [str, int, float, int, str, str, str, None]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_dialog"
+        # 将数据插入数据库
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL(
+                "COPY {} (id, day, t, type, speaker, content, created_at) FROM STDIN"
+            ).format(psycopg.sql.Identifier(table_name))
+            _rows: list[Any] = []
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
+
+    @lock_decorator
+    async def async_write_status(self, rows: list[tuple]):
+        _tuple_types = [str, int, float, float, float, int, str, str, None]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_status"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL(
+                "COPY {} (id, day, t, lng, lat, parent_id, action, status, created_at) FROM STDIN"
+            ).format(psycopg.sql.Identifier(table_name))
+            _rows: list[Any] = []
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
+
+    @lock_decorator
+    async def async_write_profile(self, rows: list[tuple]):
+        _tuple_types = [str, str, str]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_profile"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL("COPY {} (id, name, profile) FROM STDIN").format(
+                psycopg.sql.Identifier(table_name)
+            )
+            _rows: list[Any] = []
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
+
+    @lock_decorator
+    async def async_write_survey(self, rows: list[tuple]):
+        _tuple_types = [str, int, float, str, str, None]
+        table_name = f"socialcity_{self.exp_id.replace('-', '_')}_agent_survey"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            copy_sql = psycopg.sql.SQL(
+                "COPY {} (id, day, t, survey_id, result, created_at) FROM STDIN"
+            ).format(psycopg.sql.Identifier(table_name))
+            _rows: list[Any] = []
+            async with aconn.cursor() as cur:
+                async with cur.copy(copy_sql) as copy:
+                    for row in rows:
+                        _row = [
+                            _type(r) if _type is not None else r
+                            for (_type, r) in zip(_tuple_types, row)
+                        ]
+                        await copy.write_row(_row)
+                        _rows.append(_row)
+            logger.debug(f"table:{table_name} sql: {copy_sql} values: {_rows}")
+
+    @lock_decorator
+    async def async_update_exp_info(self, exp_info: dict[str, Any]):
+        # timestamp不做类型转换
+        table_name = f"socialcity_experiment"
+        async with await psycopg.AsyncConnection.connect(self._dsn) as aconn:
+            async with aconn.cursor(row_factory=dict_row) as cur:
+                exec_str = "SELECT * FROM {table_name} WHERE id=%s".format(
+                    table_name=table_name
+                ), (self.exp_id,)
+                await cur.execute(
+                    "SELECT * FROM {table_name} WHERE id=%s".format(
+                        table_name=table_name
+                    ),
+                    (self.exp_id,),
+                )  # type:ignore
+                logger.debug(f"table:{table_name} sql: {exec_str}")
+                record_exists = await cur.fetchall()
+                if record_exists:
+                    # UPDATE
+                    columns = ", ".join(
+                        f"{key} = %s" for key, _ in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                    )
+                    update_sql = psycopg.sql.SQL(
+                        f"UPDATE {{}} SET {columns} WHERE id='{self.exp_id}'"  # type:ignore
+                    ).format(psycopg.sql.Identifier(table_name))
+                    params = [
+                        _type(exp_info[key]) if _type is not None else exp_info[key]
+                        for key, _type in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                    ]
+                    logger.debug(
+                        f"table:{table_name} sql: {update_sql} values: {params}"
+                    )
+                    await cur.execute(update_sql, params)
+                else:
+                    # INSERT
+                    keys = ", ".join(
+                        key for key, _ in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                    )
+                    placeholders = ", ".join(
+                        ["%s"] * len(TO_UPDATE_EXP_INFO_KEYS_AND_TYPES)
+                    )
+                    insert_sql = psycopg.sql.SQL(
+                        f"INSERT INTO {{}} ({keys}) VALUES ({placeholders})"  # type:ignore
+                    ).format(psycopg.sql.Identifier(table_name))
+                    params = [
+                        _type(exp_info[key]) if _type is not None else exp_info[key]
+                        for key, _type in TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+                    ]
+                    logger.debug(
+                        f"table:{table_name} sql: {insert_sql} values: {params}"
+                    )
+                    await cur.execute(insert_sql, params)
+                await aconn.commit()

pycityagent/survey/__init__.py

@@ -0,0 +1,4 @@
+from .models import QuestionType, Question, Survey
+from .manager import SurveyManager
+
+__all__ = ["QuestionType", "Question", "Survey", "SurveyManager"]

pycityagent/survey/manager.py

@@ -0,0 +1,54 @@
+import json
+import uuid
+from datetime import datetime
+from typing import Optional
+
+from .models import Page, Question, QuestionType, Survey
+
+
+class SurveyManager:
+    def __init__(self):
+        self._surveys: dict[str, Survey] = {}
+
+    def create_survey(self, title: str, description: str, pages: list[dict]) -> Survey:
+        """创建新问卷"""
+        survey_id = uuid.uuid4()
+
+        # 转换页面和问题数据
+        survey_pages = []
+        for page_data in pages:
+            questions = []
+            for q in page_data["elements"]:
+                question = Question(
+                    name=q["name"],
+                    title=q["title"],
+                    type=QuestionType(q["type"]),
+                    required=q.get("required", True),
+                    choices=q.get("choices", []),
+                    columns=q.get("columns", []),
+                    rows=q.get("rows", []),
+                    min_rating=q.get("min_rating", 1),
+                    max_rating=q.get("max_rating", 5),
+                )
+                questions.append(question)
+
+            page = Page(name=page_data["name"], elements=questions)
+            survey_pages.append(page)
+
+        survey = Survey(
+            id=survey_id,
+            title=title,
+            description=description,
+            pages=survey_pages,
+        )
+
+        self._surveys[str(survey_id)] = survey
+        return survey
+
+    def get_survey(self, survey_id: str) -> Optional[Survey]:
+        """获取指定问卷"""
+        return self._surveys.get(survey_id)
+
+    def get_all_surveys(self) -> list[Survey]:
+        """获取所有问卷"""
+        return list(self._surveys.values())

pycityagent/survey/models.py

@@ -0,0 +1,120 @@
+import json
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import Enum
+from typing import Any
+
+
+class QuestionType(Enum):
+    TEXT = "text"
+    RADIO = "radiogroup"
+    CHECKBOX = "checkbox"
+    BOOLEAN = "boolean"
+    RATING = "rating"
+    MATRIX = "matrix"
+
+
+@dataclass
+class Question:
+    name: str
+    title: str
+    type: QuestionType
+    choices: list[str] = field(default_factory=list)
+    columns: list[str] = field(default_factory=list)
+    rows: list[str] = field(default_factory=list)
+    required: bool = True
+    min_rating: int = 1
+    max_rating: int = 5
+
+    def to_dict(self) -> dict:
+        base_dict: dict[str, Any] = {
+            "type": self.type.value,
+            "name": self.name,
+            "title": self.title,
+        }
+
+        if self.type in [QuestionType.RADIO, QuestionType.CHECKBOX]:
+            base_dict["choices"] = self.choices
+        elif self.type == QuestionType.MATRIX:
+            base_dict["columns"] = self.columns
+            base_dict["rows"] = self.rows
+        elif self.type == QuestionType.RATING:
+            base_dict["min_rating"] = self.min_rating
+            base_dict["max_rating"] = self.max_rating
+
+        return base_dict
+
+
+@dataclass
+class Page:
+    name: str
+    elements: list[Question]
+
+    def to_dict(self) -> dict:
+        return {"name": self.name, "elements": [q.to_dict() for q in self.elements]}
+
+
+@dataclass
+class Survey:
+    id: uuid.UUID
+    title: str
+    description: str
+    pages: list[Page]
+    responses: dict[str, dict] = field(default_factory=dict)
+    created_at: datetime = field(default_factory=datetime.now)
+
+    def to_dict(self) -> dict:
+        return {
+            "id": str(self.id),
+            "title": self.title,
+            "description": self.description,
+            "pages": [p.to_dict() for p in self.pages],
+            "response_count": len(self.responses),
+        }
+
+    def to_json(self) -> str:
+        """Convert the survey to a JSON string for MQTT transmission"""
+        survey_dict = {
+            "id": str(self.id),
+            "title": self.title,
+            "description": self.description,
+            "pages": [p.to_dict() for p in self.pages],
+            "responses": self.responses,
+            "created_at": self.created_at.isoformat(),
+        }
+        return json.dumps(survey_dict)
+
+    @classmethod
+    def from_json(cls, json_str: str) -> "Survey":
+        """Create a Survey instance from a JSON string"""
+        data = json.loads(json_str)
+        pages = [
+            Page(
+                name=p["name"],
+                elements=[
+                    Question(
+                        name=q["name"],
+                        title=q["title"],
+                        type=QuestionType(q["type"]),
+                        required=q.get("required", True),
+                        choices=q.get("choices", []),
+                        columns=q.get("columns", []),
+                        rows=q.get("rows", []),
+                        min_rating=q.get("min_rating", 1),
+                        max_rating=q.get("max_rating", 5),
+                    )
+                    for q in p["elements"]
+                ],
+            )
+            for p in data["pages"]
+        ]
+
+        return cls(
+            id=uuid.UUID(data["id"]),
+            title=data["title"],
+            description=data["description"],
+            pages=pages,
+            responses=data.get("responses", {}),
+            created_at=datetime.fromisoformat(data["created_at"]),
+        )

pycityagent/utils/__init__.py

@@ -0,0 +1,11 @@
+from .avro_schema import (DIALOG_SCHEMA, INSTITUTION_STATUS_SCHEMA,
+                          PROFILE_SCHEMA, STATUS_SCHEMA, SURVEY_SCHEMA)
+from .pg_query import PGSQL_DICT, TO_UPDATE_EXP_INFO_KEYS_AND_TYPES
+from .survey_util import process_survey_for_llm
+
+__all__ = [
+    "PROFILE_SCHEMA", "DIALOG_SCHEMA", "STATUS_SCHEMA", "SURVEY_SCHEMA", "INSTITUTION_STATUS_SCHEMA",
+    "process_survey_for_llm",
+    "TO_UPDATE_EXP_INFO_KEYS_AND_TYPES",
+    "PGSQL_DICT",
+]

pycityagent/utils/avro_schema.py

@@ -0,0 +1,109 @@
+PROFILE_SCHEMA = {
+    "doc": "Agent属性",
+    "name": "AgentProfile",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "name", "type": "string"},
+        {"name": "gender", "type": "string"},
+        {"name": "age", "type": "float"},
+        {"name": "education", "type": "string"},
+        {"name": "skill", "type": "string"},
+        {"name": "occupation", "type": "string"},
+        {"name": "family_consumption", "type": "string"},
+        {"name": "consumption", "type": "string"},
+        {"name": "personality", "type": "string"},
+        {"name": "income", "type": "string"},
+        {"name": "currency", "type": "float"},
+        {"name": "residence", "type": "string"},
+        {"name": "race", "type": "string"},
+        {"name": "religion", "type": "string"},
+        {"name": "marital_status", "type": "string"},
+    ],
+}
+
+DIALOG_SCHEMA = {
+    "doc": "Agent对话",
+    "name": "AgentDialog",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "type", "type": "int"},
+        {"name": "speaker", "type": "string"},
+        {"name": "content", "type": "string"},
+        {
+            "name": "created_at",
+            "type": {"type": "long", "logicalType": "timestamp-millis"},
+        },
+    ],
+}
+
+STATUS_SCHEMA = {
+    "doc": "Agent状态",
+    "name": "AgentStatus",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "lng", "type": "double"},
+        {"name": "lat", "type": "double"},
+        {"name": "parent_id", "type": "int"},
+        {"name": "action", "type": "string"},
+        {"name": "hungry", "type": "float"},
+        {"name": "tired", "type": "float"},
+        {"name": "safe", "type": "float"},
+        {"name": "social", "type": "float"},
+        {
+            "name": "created_at",
+            "type": {"type": "long", "logicalType": "timestamp-millis"},
+        },
+    ],
+}
+
+INSTITUTION_STATUS_SCHEMA = {
+    "doc": "Institution状态",
+    "name": "InstitutionStatus",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "type", "type": "int"},
+        {"name": "nominal_gdp", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "real_gdp", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "unemployment", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "wages", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "prices", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "inventory", "type": ["int", "null"]},
+        {"name": "price", "type": ["float", "null"]},
+        {"name": "interest_rate", "type": ["float", "null"]},
+        {"name": "bracket_cutoffs", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "bracket_rates", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+        {"name": "employees", "type": {"type": "array", "items": ["float", "int", "string", "null"]}},
+    ],
+}
+
+SURVEY_SCHEMA = {
+    "doc": "Agent问卷",
+    "name": "AgentSurvey",
+    "namespace": "com.socialcity",
+    "type": "record",
+    "fields": [
+        {"name": "id", "type": "string"},  # uuid as string
+        {"name": "day", "type": "int"},
+        {"name": "t", "type": "float"},
+        {"name": "survey_id", "type": "string"},
+        {"name": "result", "type": "string"},
+        {
+            "name": "created_at",
+            "type": {"type": "long", "logicalType": "timestamp-millis"},
+        },
+    ],
+}

pycityagent/utils/decorators.py

@@ -0,0 +1,99 @@
+import time
+import functools
+import inspect
+
+CALLING_STRING = 'function: `{func_name}` in "{file_path}", line {line_number}, arguments: `{arguments}` start time: `{start_time}` end time: `{end_time}` output: `{output}`'
+
+__all__ = [
+    "record_call_aio",
+    "record_call",
+    "lock_decorator",
+]
+
+
+def record_call_aio(record_function_calling: bool = True):
+    """
+    Decorator to log the async function call details if `record_function_calling` is True.
+    """
+
+    def decorator(func):
+        async def wrapper(*args, **kwargs):
+            cur_frame = inspect.currentframe()
+            assert cur_frame is not None
+            frame = cur_frame.f_back
+            assert frame is not None
+            line_number = frame.f_lineno
+            file_path = frame.f_code.co_filename
+            args_repr = [repr(a) for a in args]
+            kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
+            signature = ", ".join(args_repr + kwargs_repr)
+            start_time = time.time()
+            result = await func(*args, **kwargs)
+            end_time = time.time()
+            if record_function_calling:
+                print(
+                    CALLING_STRING.format(
+                        func_name=func,
+                        line_number=line_number,
+                        file_path=file_path,
+                        arguments=signature,
+                        start_time=start_time,
+                        end_time=end_time,
+                        output=result,
+                    )
+                )
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def record_call(record_function_calling: bool = True):
+    """
+    Decorator to log the function call details if `record_function_calling` is True.
+    """
+
+    def decorator(func):
+        def wrapper(*args, **kwargs):
+            cur_frame = inspect.currentframe()
+            assert cur_frame is not None
+            frame = cur_frame.f_back
+            assert frame is not None
+            line_number = frame.f_lineno
+            file_path = frame.f_code.co_filename
+            args_repr = [repr(a) for a in args]
+            kwargs_repr = [f"{k}={v!r}" for k, v in kwargs.items()]
+            signature = ", ".join(args_repr + kwargs_repr)
+            start_time = time.time()
+            result = func(*args, **kwargs)
+            end_time = time.time()
+            if record_function_calling:
+                print(
+                    CALLING_STRING.format(
+                        func_name=func,
+                        line_number=line_number,
+                        file_path=file_path,
+                        arguments=signature,
+                        start_time=start_time,
+                        end_time=end_time,
+                        output=result,
+                    )
+                )
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def lock_decorator(func):
+    async def wrapper(self, *args, **kwargs):
+        lock = self._lock
+        await lock.acquire()
+        try:
+            return await func(self, *args, **kwargs)
+        finally:
+            lock.release()
+
+    return wrapper

pycityagent/utils/parsers/__init__.py

@@ -0,0 +1,13 @@
+"""Model response parser module."""
+
+from .parser_base import ParserBase
+from .json_parser import JsonDictParser, JsonObjectParser
+from .code_block_parser import CodeBlockParser
+
+
+__all__ = [
+    "ParserBase",
+    "JsonDictParser",
+    "JsonObjectParser",
+    "CodeBlockParser",
+]

pycityagent/utils/parsers/code_block_parser.py

@@ -0,0 +1,37 @@
+import logging
+from typing import Any
+
+from .parser_base import ParserBase
+
+
+class CodeBlockParser(ParserBase):
+    """A parser that extracts specific objects from a response string enclosed within specific tags.
+
+    Attributes:
+        tag_start (str): The start tag used to identify the beginning of the object.
+        tag_end (str): The end tag used to identify the end of the object.
+    """
+
+    tag_start = "```{language_name}"
+    tag_end = "```"
+
+    def __init__(self, language_name: str) -> None:
+        """Initialize the CodeBlockParser with default tags."""
+        super().__init__()
+        self.language_name = language_name
+
+    def parse(self, response: str) -> str:
+        """Parse the response string to extract and return a str.
+
+        Parameters:
+            response (str): The response string containing the specified language object.
+
+        Returns:
+            str: The parsed `str` object.
+        """
+        extract_text = self._extract_text_within_tags(
+            response=response,
+            tag_start=self.tag_start.format(language_name=self.language_name),
+            tag_end=self.tag_end,
+        )
+        return extract_text