ForcePilot/backend/test/api/test_attachment_and_agent_state.py

"""
测试附件上传和 agent state 获取的 API 脚本

使用方式:
    cd /home/zwj/workspace/Yuxi-Know
    docker compose exec api uv run python test/api/test_attachment_and_agent_state.py

或者本地运行:
    python test/api/test_attachment_and_agent_state.py
"""

import asyncio
import contextlib
import os
import sys
import uuid
from pathlib import Path

import httpx
from dotenv import load_dotenv

# 添加项目根目录到 Python 路径
PROJECT_ROOT = Path(__file__).parent.parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

load_dotenv(PROJECT_ROOT / ".env")


# API 配置
API_BASE_URL = os.getenv("API_BASE_URL", "http://localhost:5050")

# 测试账户配置
USERNAME = os.getenv("YUXI_SUPER_ADMIN_NAME", "zwj")
PASSWORD = os.getenv("YUXI_SUPER_ADMIN_PASSWORD", "zwj12138")

# 默认 Agent ID (需要根据实际情况修改，使用类名)
DEFAULT_AGENT_ID = "ChatbotAgent"


class APITester:
    def __init__(self, base_url: str, username: str, password: str):
        self.base_url = base_url
        self.username = username
        self.password = password
        self.token: str | None = None
        self.user_id: str | None = None
        self.headers: dict | None = None

    @contextlib.asynccontextmanager
    async def _client(self, timeout: float = 30.0):
        """获取 HTTP 客户端"""
        client = httpx.AsyncClient(timeout=timeout)
        try:
            yield client
        finally:
            await client.aclose()

    async def login(self) -> bool:
        """登录获取 token"""
        print(f"\n{'=' * 60}")
        print(f"1. 正在登录: {self.username}")
        print(f"{'=' * 60}")

        async with self._client() as client:
            response = await client.post(
                f"{self.base_url}/api/auth/token",
                data={"username": self.username, "password": self.password},
            )

        if response.status_code == 200:
            data = response.json()
            self.token = data.get("access_token")
            self.user_id = str(data.get("user_id"))
            self.headers = {"Authorization": f"Bearer {self.token}"}
            print(f"   ✓ 登录成功! user_id: {self.user_id}")
            print(f"   ✓ Token: {self.token[:50]}...")
            return True
        else:
            print(f"   ✗ 登录失败: {response.status_code} - {response.text}")
            return False

    async def create_thread(self, agent_id: str) -> str:
        """创建对话线程"""
        print(f"\n{'=' * 60}")
        print(f"2. 创建对话线程 (agent_id: {agent_id})")
        print(f"{'=' * 60}")

        async with self._client() as client:
            response = await client.post(
                f"{self.base_url}/api/chat/thread",
                json={"agent_id": agent_id, "title": "API 测试对话"},
                headers=self.headers,
            )

        if response.status_code == 200:
            thread = response.json()
            thread_id = thread.get("id")
            print(f"   ✓ 创建成功! thread_id: {thread_id}")
            return thread_id
        else:
            print(f"   ✗ 创建失败: {response.status_code} - {response.text}")
            return ""

    async def upload_attachment(self, thread_id: str, file_path: str) -> dict | None:
        """上传附件"""
        print(f"\n{'=' * 60}")
        print(f"3. 上传附件: {file_path}")
        print(f"{'=' * 60}")

        if not os.path.exists(file_path):
            print(f"   ✗ 文件不存在: {file_path}")
            return None

        async with self._client(timeout=60.0) as client:
            with open(file_path, "rb") as f:
                files = {"file": (os.path.basename(file_path), f)}
                response = await client.post(
                    f"{self.base_url}/api/chat/thread/{thread_id}/attachments",
                    files=files,
                    headers=self.headers,
                )

        if response.status_code == 200:
            attachment = response.json()
            print("   ✓ 上传成功!")
            print(f"      file_id: {attachment.get('file_id')}")
            print(f"      file_name: {attachment.get('file_name')}")
            print(f"      status: {attachment.get('status')}")
            return attachment
        else:
            print(f"   ✗ 上传失败: {response.status_code} - {response.text}")
            return None

    async def list_attachments(self, thread_id: str) -> list[dict]:
        """列出附件"""
        print(f"\n{'=' * 60}")
        print(f"4. 列出附件 (thread_id: {thread_id})")
        print(f"{'=' * 60}")

        async with self._client() as client:
            response = await client.get(
                f"{self.base_url}/api/chat/thread/{thread_id}/attachments",
                headers=self.headers,
            )

        if response.status_code == 200:
            data = response.json()
            attachments = data.get("attachments", [])
            print(f"   ✓ 获取到 {len(attachments)} 个附件:")
            for att in attachments:
                print(f"      - {att.get('file_name')}: {att.get('status')}")
            return attachments
        else:
            print(f"   ✗ 获取失败: {response.status_code} - {response.text}")
            return []

    async def get_agent_state(self, agent_id: str, thread_id: str) -> dict | None:
        """获取 agent state"""
        print(f"\n{'=' * 60}")
        print(f"5. 获取 Agent State (agent_id: {agent_id}, thread_id: {thread_id})")
        print(f"{'=' * 60}")

        async with self._client() as client:
            response = await client.get(
                f"{self.base_url}/api/chat/agent/{agent_id}/state",
                params={"thread_id": thread_id},
                headers=self.headers,
            )

        if response.status_code == 200:
            state = response.json()
            agent_state = state.get("agent_state", {})
            print("   ✓ 获取成功!")
            print(f"      files: {len(agent_state.get('files', {}))} 个")
            print(f"      todos: {len(agent_state.get('todos', []))} 个")
            if agent_state.get("files"):
                print("      文件列表:")
                for path in agent_state["files"]:
                    file_info = agent_state["files"][path]
                    print(f"         - {path}: {len(file_info.get('content', []))} 行")
            return state
        else:
            print(f"   ✗ 获取失败: {response.status_code} - {response.text}")
            return None

    async def get_agent_config_id(self, agent_id: str) -> int | None:
        """获取指定 Agent 的一个配置 ID"""
        async with self._client() as client:
            response = await client.get(
                f"{self.base_url}/api/chat/agent/{agent_id}/configs",
                headers=self.headers,
            )

        if response.status_code != 200:
            print(f"   ✗ 获取 Agent 配置失败: {response.status_code} - {response.text}")
            return None

        configs = response.json().get("configs", [])
        if not configs:
            print("   ✗ 当前 Agent 没有可用配置")
            return None

        return configs[0].get("id")

    async def send_chat_message(self, agent_id: str, thread_id: str, query: str) -> bool:
        """发送聊天消息（流式）"""
        print(f"\n{'=' * 60}")
        print("6. 发送聊天消息")
        print(f"{'=' * 60}")
        print(f"   Query: {query}")
        print(f"   Thread ID: {thread_id}")

        agent_config_id = await self.get_agent_config_id(agent_id)
        if not agent_config_id:
            return False

        async with self._client(timeout=120.0) as client:
            async with client.stream(
                "POST",
                f"{self.base_url}/api/chat/agent",
                json={
                    "query": query,
                    "agent_config_id": agent_config_id,
                    "thread_id": thread_id,
                },
                headers=self.headers,
            ) as response:
                print(f"\n   响应状态: {response.status_code}")
                print("   响应内容:")
                async for chunk in response.aiter_lines():
                    if chunk:
                        print(f"      {chunk[:150]}...")

                return response.status_code == 200


async def main():
    """主测试流程"""
    print("\n" + "=" * 60)
    print("   附件上传与 Agent State API 测试")
    print("=" * 60)
    print(f"\nAPI 地址: {API_BASE_URL}")
    print(f"测试账户: {USERNAME}")

    tester = APITester(API_BASE_URL, USERNAME, PASSWORD)

    # 1. 登录
    if not await tester.login():
        print("\n!!! 登录失败，测试终止 !!!")
        return

    # 2. 创建线程（使用默认 agent_id）
    agent_id = DEFAULT_AGENT_ID
    thread_id = await tester.create_thread(agent_id)
    if not thread_id:
        print("\n!!! 创建线程失败，测试终止 !!!")
        return

    # 3. 创建测试文件
    test_content = """# 测试文档

这是一个用于 API 测试的 Markdown 文件。

## 主要内容

- 第一点
- 第二点
- 第三点

```python
def hello():
    print("Hello, World!")
```
"""
    test_file_path = f"/tmp/test_attachment_{uuid.uuid4().hex[:8]}.md"
    with open(test_file_path, "w", encoding="utf-8") as f:
        f.write(test_content)
    print(f"\n   测试文件已创建: {test_file_path}")

    # 4. 上传附件
    attachment = await tester.upload_attachment(thread_id, test_file_path)

    # 5. 列出附件
    await tester.list_attachments(thread_id)

    # 6. 获取 agent state (验证附件是否在 state 中)
    if attachment:
        print("\n   等待后端处理...")
        await asyncio.sleep(2)
        await tester.get_agent_state(agent_id, thread_id)

    # 7. 发送聊天消息测试
    await tester.send_chat_message(agent_id, thread_id, "你好，请简单介绍一下你自己。")

    # 8. 再次获取 agent state (验证 todos 等状态)
    await asyncio.sleep(1)
    await tester.get_agent_state(agent_id, thread_id)

    # 清理测试文件
    if os.path.exists(test_file_path):
        os.remove(test_file_path)
        print(f"\n   测试文件已清理: {test_file_path}")

    print(f"\n{'=' * 60}")
    print("   测试完成!")
    print(f"{'=' * 60}\n")


if __name__ == "__main__":
    asyncio.run(main())