ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter9/gaia-experience/AWorld/aworlddistributed/aworldspace/db/db.py

@@ -0,0 +1,210 @@
+from abc import ABC, abstractmethod
+from datetime import datetime
+from typing import Optional
+
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+
+from base import AworldTask, AworldTaskResult
+from aworldspace.db.models import (
+    Base, AworldTaskModel, AworldTaskResultModel,
+    orm_to_pydantic_task, pydantic_to_orm_task,
+    orm_to_pydantic_result, pydantic_to_orm_result
+)
+
+
+class AworldTaskDB(ABC):
+
+    @abstractmethod
+    async def query_task_by_id(self, task_id: str) -> AworldTask:
+        pass
+
+    @abstractmethod
+    async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
+        pass
+
+    @abstractmethod
+    async def insert_task(self, task: AworldTask):
+        pass
+
+    @abstractmethod
+    async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
+        pass
+
+    @abstractmethod
+    async def update_task(self, task: AworldTask):
+        pass
+
+    @abstractmethod
+    async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
+        pass
+
+    @abstractmethod
+    async def save_task_result(self, result: AworldTaskResult):
+        pass
+
+
+class SqliteTaskDB(AworldTaskDB):
+    def __init__(self, db_path: str):
+        self.engine = create_engine(db_path, echo=False, future=True)
+        Base.metadata.create_all(self.engine)
+        self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
+
+    async def query_task_by_id(self, task_id: str) -> Optional[AworldTask]:
+        with self.Session() as session:
+            orm_task = session.query(AworldTaskModel).filter_by(task_id=task_id).first()
+            return orm_to_pydantic_task(orm_task) if orm_task else None
+
+    async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
+        with self.Session() as session:
+            orm_result = (
+                session.query(AworldTaskResultModel)
+                .filter_by(task_id=task_id)
+                .order_by(AworldTaskResultModel.created_at.desc())
+                .first()
+            )
+            return orm_to_pydantic_result(orm_result) if orm_result else None
+
+    async def insert_task(self, task: AworldTask):
+        with self.Session() as session:
+            orm_task = pydantic_to_orm_task(task)
+            session.add(orm_task)
+            session.commit()
+
+    async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
+        with self.Session() as session:
+            orm_tasks = (
+                session.query(AworldTaskModel)
+                .filter_by(status=status)
+                .limit(nums)
+                .all()
+            )
+            return [orm_to_pydantic_task(t) for t in orm_tasks]
+
+    async def update_task(self, task: AworldTask):
+        with self.Session() as session:
+            orm_task = session.query(AworldTaskModel).filter_by(task_id=task.task_id).first()
+            if orm_task:
+                for k, v in task.model_dump().items():
+                    setattr(orm_task, k, v)
+                orm_task.updated_at = datetime.utcnow()
+                session.commit()
+
+    async def save_task_result(self, result: AworldTaskResult):
+        with self.Session() as session:
+            orm_task = pydantic_to_orm_result(result)
+            session.add(orm_task)
+            session.commit()
+
+    async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
+        with self.Session() as session:
+            query = session.query(AworldTaskModel)
+            
+            # Handle special filters for time ranges
+            start_time = filter.pop('start_time', None)
+            end_time = filter.pop('end_time', None)
+            
+            # Apply regular filters
+            for k, v in filter.items():
+                if hasattr(AworldTaskModel, k):
+                    query = query.filter(getattr(AworldTaskModel, k) == v)
+            
+            # Apply time range filters
+            if start_time:
+                query = query.filter(AworldTaskModel.created_at >= start_time)
+            if end_time:
+                query = query.filter(AworldTaskModel.created_at <= end_time)
+            
+            total = query.count()
+            orm_tasks = query.offset((page_num - 1) * page_size).limit(page_size).all()
+            items = [orm_to_pydantic_task(t) for t in orm_tasks]
+            return {
+                "total": total,
+                "page_num": page_num,
+                "page_size": page_size,
+                "items": items
+            }
+
+
+class PostgresTaskDB(AworldTaskDB):
+    def __init__(self, db_url: str):
+        # db_url example: 'postgresql+psycopg2://user:password@host:port/dbname'
+        self.engine = create_engine(db_url, echo=False, future=True)
+        Base.metadata.create_all(self.engine)
+        self.Session = sessionmaker(bind=self.engine, expire_on_commit=False)
+
+    async def query_task_by_id(self, task_id: str) -> Optional[AworldTask]:
+        with self.Session() as session:
+            orm_task = session.query(AworldTaskModel).filter_by(task_id=task_id).first()
+            return orm_to_pydantic_task(orm_task) if orm_task else None
+
+    async def query_latest_task_result_by_id(self, task_id: str) -> Optional[AworldTaskResult]:
+        with self.Session() as session:
+            orm_result = (
+                session.query(AworldTaskResultModel)
+                .filter_by(task_id=task_id)
+                .order_by(AworldTaskResultModel.created_at.desc())
+                .first()
+            )
+            return orm_to_pydantic_result(orm_result) if orm_result else None
+
+    async def insert_task(self, task: AworldTask):
+        with self.Session() as session:
+            orm_task = pydantic_to_orm_task(task)
+            session.add(orm_task)
+            session.commit()
+
+    async def query_tasks_by_status(self, status: str, nums: int) -> list[AworldTask]:
+        with self.Session() as session:
+            orm_tasks = (
+                session.query(AworldTaskModel)
+                .filter_by(status=status)
+                .limit(nums)
+                .all()
+            )
+            return [orm_to_pydantic_task(t) for t in orm_tasks]
+
+    async def update_task(self, task: AworldTask):
+        with self.Session() as session:
+            orm_task = session.query(AworldTaskModel).filter_by(task_id=task.task_id).first()
+            if orm_task:
+                for k, v in task.model_dump().items():
+                    setattr(orm_task, k, v)
+                orm_task.updated_at = datetime.utcnow()
+                session.commit()
+
+    async def save_task_result(self, result: AworldTaskResult):
+        with self.Session() as session:
+            orm_task = pydantic_to_orm_result(result)
+            session.add(orm_task)
+            session.commit()
+
+    async def page_query_tasks(self, filter: dict, page_size: int, page_num: int) -> dict:
+        with self.Session() as session:
+            query = session.query(AworldTaskModel)
+            
+            # Handle special filters for time ranges
+            start_time = filter.pop('start_time', None)
+            end_time = filter.pop('end_time', None)
+            
+            # Apply regular filters
+            for k, v in filter.items():
+                if hasattr(AworldTaskModel, k):
+                    query = query.filter(getattr(AworldTaskModel, k) == v)
+            
+            # Apply time range filters
+            if start_time:
+                query = query.filter(AworldTaskModel.created_at >= start_time)
+            if end_time:
+                query = query.filter(AworldTaskModel.created_at <= end_time)
+            
+            total = query.count()
+            orm_tasks = query.offset((page_num - 1) * page_size).limit(page_size).all()
+            items = [orm_to_pydantic_task(t) for t in orm_tasks]
+            return {
+                "total": total,
+                "page_num": page_num,
+                "page_size": page_size,
+                "items": items
+            }
+

chapter9/gaia-experience/AWorld/aworlddistributed/aworldspace/db/models.py

@@ -0,0 +1,66 @@
+from sqlalchemy import Column, String, Integer, Text, DateTime, JSON, create_engine
+from sqlalchemy.orm import declarative_base
+from datetime import datetime
+from typing import Optional
+from base import AworldTask, AworldTaskResult
+
+Base = declarative_base()
+
+
+class AworldTaskModel(Base):
+    __tablename__ = 'aworld_tasks'
+
+    task_id = Column(String, primary_key=True)
+    agent_id = Column(String)
+    agent_input = Column(Text)
+    session_id = Column(String)
+    user_id = Column(String)
+    llm_provider = Column(String)
+    llm_model_name = Column(String)
+    llm_api_key = Column(String)
+    llm_base_url = Column(String)
+    llm_custom_input = Column(Text)
+    task_system_prompt = Column(Text)
+    mcp_servers = Column(JSON)
+    node_id = Column(String)
+    client_id = Column(String)
+    status = Column(String, default='INIT')
+    history_messages = Column(Integer, default=100)
+    max_steps = Column(Integer, default=100)
+    max_retries = Column(Integer, default=5)
+    ext_info = Column(JSON, default=dict)
+    created_at = Column(DateTime, default=datetime.utcnow)
+    updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
+
+
+class AworldTaskResultModel(Base):
+    __tablename__ = 'aworld_tasks_results'
+
+    task_result_id = Column(Integer, primary_key=True, autoincrement=True)
+    task_id = Column(String)
+    server_host = Column(String)
+    data = Column(JSON)
+    created_at = Column(DateTime, default=datetime.utcnow)
+
+
+def orm_to_pydantic_task(orm_obj: AworldTaskModel) -> AworldTask:
+    return AworldTask(**{c.name: getattr(orm_obj, c.name) for c in orm_obj.__table__.columns})
+
+
+def pydantic_to_orm_task(pydantic_obj: AworldTask) -> AworldTaskModel:
+    return AworldTaskModel(**pydantic_obj.model_dump())
+
+
+def orm_to_pydantic_result(orm_obj: AworldTaskResultModel) -> AworldTaskResult:
+    return AworldTaskResult(
+        server_host=orm_obj.server_host,
+        data=orm_obj.data
+    )
+
+
+def pydantic_to_orm_result(pydantic_obj: AworldTaskResult) -> AworldTaskResultModel:
+    return AworldTaskResultModel(
+        task_id=pydantic_obj.task.task_id if pydantic_obj.task else None,
+        server_host=pydantic_obj.server_host,
+        data=pydantic_obj.data
+    )