agent_management/agent_templates/test_client.py

#!/usr/bin/env python3
"""
Azure Blob Storage Agent 客户端示例
演示如何使用 Python 调用 agent API
"""
import requests
import json
import os
import sys

# Agent 配置
AGENT_BASE_URL = os.getenv("AGENT_URL", "http://localhost:8080")

class AzureBlobAgentClient:
    """Azure Blob Storage Agent 客户端"""
    
    def __init__(self, base_url: str = AGENT_BASE_URL):
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
        self.connected = False
    
    def health_check(self) -> dict:
        """健康检查"""
        response = self.session.get(f"{self.base_url}/health")
        response.raise_for_status()
        return response.json()
    
    def connect(self, connection_string: str) -> dict:
        """连接到 Azure Storage"""
        response = self.session.post(
            f"{self.base_url}/connect",
            json={"connection_string": connection_string}
        )
        response.raise_for_status()
        result = response.json()
        self.connected = True
        return result
    
    def query(self, query_text: str, container_name: str = None) -> dict:
        """执行自然语言查询"""
        if not self.connected:
            raise Exception("未连接到 Azure Storage，请先调用 connect()")
        
        payload = {"query": query_text}
        if container_name:
            payload["container_name"] = container_name
        
        response = self.session.post(
            f"{self.base_url}/query",
            json=payload
        )
        response.raise_for_status()
        return response.json()
    
    def get_info(self) -> dict:
        """获取 agent 信息"""
        response = self.session.get(f"{self.base_url}/")
        response.raise_for_status()
        return response.json()


def print_response(title: str, response: dict):
    """格式化打印响应"""
    print(f"\n{'='*60}")
    print(f"📋 {title}")
    print('='*60)
    print(json.dumps(response, indent=2, ensure_ascii=False))


def main():
    """主函数"""
    print("🚀 Azure Blob Storage Agent 客户端")
    print(f"连接到: {AGENT_BASE_URL}\n")
    
    # 创建客户端
    client = AzureBlobAgentClient()
    
    try:
        # 1. 健康检查
        print("1️⃣  执行健康检查...")
        health = client.health_check()
        print_response("健康检查", health)
        
        # 2. 获取 agent 信息
        print("\n2️⃣  获取 Agent 信息...")
        info = client.get_info()
        print_response("Agent 信息", info)
        
        # 3. 连接到 Azure Storage
        print("\n3️⃣  连接到 Azure Storage...")
        
        # 从环境变量获取连接字符串
        connection_string = os.getenv("AZURE_STORAGE_CONNECTION_STRING")
        
        if not connection_string:
            print("\n⚠️  未设置 AZURE_STORAGE_CONNECTION_STRING 环境变量")
            print("请输入 Azure Storage 连接字符串:")
            connection_string = input().strip()
        
        if not connection_string:
            print("❌ 未提供连接字符串，退出")
            sys.exit(1)
        
        connect_result = client.connect(connection_string)
        print_response("连接结果", connect_result)
        
        # 4. 执行查询
        print("\n4️⃣  执行自然语言查询...\n")
        
        queries = [
            "列出所有容器",
            "显示存储统计信息",
        ]
        
        for query_text in queries:
            print(f"\n💬 查询: {query_text}")
            result = client.query(query_text)
            print(f"\n✅ 答案:\n{result.get('answer', 'N/A')}")
            print(f"\n状态: {result.get('status')}")
        
        # 5. 交互式查询
        print("\n5️⃣  交互式查询")
        print("="*60)
        print("输入自然语言查询（输入 'quit' 或 'exit' 退出）:")
        print("例如:")
        print("  - 列出所有容器")
        print("  - 显示 images 容器中的文件")
        print("  - 在 documents 容器中搜索 report")
        print("  - 获取存储统计信息")
        print("="*60)
        
        while True:
            try:
                query_text = input("\n💬 > ").strip()
                
                if query_text.lower() in ['quit', 'exit', 'q']:
                    print("👋 再见!")
                    break
                
                if not query_text:
                    continue
                
                result = client.query(query_text)
                print(f"\n✅ 答案:\n{result.get('answer', 'N/A')}")
                
            except KeyboardInterrupt:
                print("\n\n👋 再见!")
                break
            except Exception as e:
                print(f"\n❌ 查询失败: {str(e)}")
    
    except requests.exceptions.ConnectionError:
        print(f"\n❌ 无法连接到 Agent: {AGENT_BASE_URL}")
        print("请确保 Agent 正在运行")
        sys.exit(1)
    except Exception as e:
        print(f"\n❌ 错误: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()