From 532dca13f4b18af7f071ef99691269bbae688d2d Mon Sep 17 00:00:00 2001 From: elipitc Date: Sun, 31 May 2026 22:33:49 +0800 Subject: [PATCH] Update sub-mode runtime model handling --- .../agents/a2a_litellm_agent/a2a_server.py | 23 +- .../agents/a2a_litellm_agent/agent.py | 358 +++++++++++++++++- .../agents/a2a_litellm_agent/config.py | 37 +- api/swarm/orchestrator.py | 110 +++++- api/swarm/router.py | 28 +- k8s_manager.py | 26 +- scripts/e2e_complete_project.py | 311 +++++++++++++++ 7 files changed, 865 insertions(+), 28 deletions(-) create mode 100755 scripts/e2e_complete_project.py diff --git a/agent_templates/agents/a2a_litellm_agent/a2a_server.py b/agent_templates/agents/a2a_litellm_agent/a2a_server.py index 644a898..33dae4d 100644 --- a/agent_templates/agents/a2a_litellm_agent/a2a_server.py +++ b/agent_templates/agents/a2a_litellm_agent/a2a_server.py @@ -18,7 +18,7 @@ from fastapi.middleware.cors import CORSMiddleware from pydantic import BaseModel, Field import structlog -from agent import LiteLLMAgent +from agent import LiteLLMAgent, ModelRequestError from config import get_config, AgentConfig, A2AConfig try: @@ -94,6 +94,7 @@ class A2ATask(BaseModel): contextId: str = Field(default_factory=lambda: uuid.uuid4().hex) status: A2ATaskStatus artifacts: Optional[list[A2AArtifact]] = None + metadata: Optional[Dict[str, Any]] = None class A2AResponse(BaseModel): @@ -392,12 +393,12 @@ class A2AAgentServer: request_id=task_id ) as ctx: ctx.add_tool("a2a_chat") - response_text = await agent.chat( + response = await agent.chat_result( message=user_text, conversation_id=context_id ) else: - response_text = await agent.chat( + response = await agent.chat_result( message=user_text, conversation_id=context_id ) @@ -411,9 +412,17 @@ class A2AAgentServer: task.artifacts = [ A2AArtifact( name="response", - parts=[A2APart(kind="text", text=response_text)] + parts=[A2APart(kind="text", text=response.get("content", ""))] ) ] + task.metadata = { + "newapi_request_id": response.get("request_id"), + "response_id": response.get("response_id"), + "model": response.get("model"), + "api_format": response.get("api_format"), + "endpoint": response.get("endpoint"), + "usage": response.get("usage") or {}, + } self.tasks[task_id] = task return JSONResponse({ @@ -425,6 +434,9 @@ class A2AAgentServer: except Exception as e: logger.error("处理消息失败", error=str(e)) task.status = A2ATaskStatus(state="failed", message=str(e)) + error_data = {} + if isinstance(e, ModelRequestError): + error_data = e.to_dict() self.tasks[task_id] = task return JSONResponse({ @@ -432,7 +444,8 @@ class A2AAgentServer: "id": request.id, "error": { "code": -32000, - "message": f"Agent error: {str(e)}" + "message": f"Agent error: {str(e)}", + "data": error_data, } }) diff --git a/agent_templates/agents/a2a_litellm_agent/agent.py b/agent_templates/agents/a2a_litellm_agent/agent.py index efdd611..b4da516 100644 --- a/agent_templates/agents/a2a_litellm_agent/agent.py +++ b/agent_templates/agents/a2a_litellm_agent/agent.py @@ -6,7 +6,7 @@ LiteLLM Agent 核心模块 import asyncio import json import uuid -from typing import AsyncGenerator, Optional, Dict, Any, List +from typing import AsyncGenerator, Optional, Dict, Any, List, Union from dataclasses import dataclass, field from datetime import datetime @@ -47,6 +47,54 @@ class Conversation: return [{"role": m.role, "content": m.content} for m in self.messages] +@dataclass +class ModelResult: + """Normalized model response metadata for Runtime accounting.""" + + content: str + usage: Dict[str, int] = field(default_factory=dict) + request_id: Optional[str] = None + response_id: Optional[str] = None + model: Optional[str] = None + api_format: str = "openai_chat" + endpoint: Optional[str] = None + + def to_dict(self) -> Dict[str, Any]: + return { + "content": self.content, + "usage": self.usage, + "request_id": self.request_id, + "response_id": self.response_id, + "model": self.model, + "api_format": self.api_format, + "endpoint": self.endpoint, + } + + +class ModelRequestError(RuntimeError): + """Model gateway error carrying request metadata for Runtime logs.""" + + def __init__( + self, + message: str, + *, + request_id: Optional[str] = None, + status_code: Optional[int] = None, + response_text: Optional[str] = None, + ): + super().__init__(message) + self.request_id = request_id + self.status_code = status_code + self.response_text = response_text + + def to_dict(self) -> Dict[str, Any]: + return { + "request_id": self.request_id, + "status_code": self.status_code, + "response_text": self.response_text, + } + + class LiteLLMAgent: """ 基于LiteLLM的Agent实现 @@ -110,6 +158,8 @@ class LiteLLMAgent: timeout=httpx.Timeout(self.llm_config.timeout), headers={ "Authorization": f"Bearer {self.llm_config.api_key}", + "x-api-key": self.llm_config.api_key, + "anthropic-version": "2023-06-01", "Content-Type": "application/json" } ) @@ -144,7 +194,7 @@ class LiteLLMAgent: message: str, conversation_id: Optional[str] = None, stream: bool = False - ) -> str | AsyncGenerator[str, None]: + ) -> Union[str, AsyncGenerator[str, None]]: """ 发送消息并获取回复 @@ -162,12 +212,73 @@ class LiteLLMAgent: conversation.add_message("user", message) if stream: + if self.llm_config.api_format == "anthropic_messages": + return self._stream_anthropic_messages_text(conversation) return self._stream_chat(conversation) - else: - return await self._simple_chat(conversation) + result = await self.chat_result_for_conversation(conversation) + return result.content + + async def chat_result( + self, + message: str, + conversation_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Return assistant text plus usage and NewAPI request metadata.""" + conversation = self.get_or_create_conversation(conversation_id) + conversation.add_message("user", message) + return (await self.chat_result_for_conversation(conversation)).to_dict() + + async def chat_result_for_conversation(self, conversation: Conversation) -> ModelResult: + """Dispatch to the configured model API format.""" + if self.llm_config.api_format == "anthropic_messages": + if self.llm_config.use_stream: + return await self._anthropic_messages_stream(conversation) + return await self._anthropic_messages(conversation) + if self.llm_config.use_stream: + return await self._openai_chat_stream_result(conversation) + return await self._simple_chat(conversation) - async def _simple_chat(self, conversation: Conversation) -> str: - """非流式对话""" + def _request_id_from_response(self, response: httpx.Response, body: Optional[Dict[str, Any]] = None) -> Optional[str]: + """Extract NewAPI/OpenAI/Anthropic request ID from headers or body.""" + for name in ( + "x-request-id", + "request-id", + "x-newapi-request-id", + "x-litellm-request-id", + "anthropic-request-id", + ): + value = response.headers.get(name) + if value: + return value + if body: + return body.get("request_id") + return None + + def _normalize_usage(self, usage: Optional[Dict[str, Any]]) -> Dict[str, int]: + usage = usage or {} + prompt_tokens = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) + completion_tokens = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) + if "input_tokens" in usage or "output_tokens" in usage: + total_tokens = prompt_tokens + completion_tokens + else: + total_tokens = int(usage.get("total_tokens") or prompt_tokens + completion_tokens) + return { + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + } + + def _raise_gateway_error(self, exc: httpx.HTTPStatusError, body: Optional[Dict[str, Any]] = None) -> None: + request_id = self._request_id_from_response(exc.response, body) + raise ModelRequestError( + f"Server error '{exc.response.status_code} {exc.response.reason_phrase}' for url '{exc.request.url}'", + request_id=request_id, + status_code=exc.response.status_code, + response_text=exc.response.text[:2000], + ) from exc + + async def _simple_chat(self, conversation: Conversation) -> ModelResult: + """非流式 OpenAI chat completions 对话""" client = await self._get_client() request_body = { @@ -184,7 +295,10 @@ class LiteLLMAgent: self.llm_config.chat_endpoint, json=request_body ) - response.raise_for_status() + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + self._raise_gateway_error(exc) result = response.json() assistant_message = result["choices"][0]["message"]["content"] @@ -192,8 +306,17 @@ class LiteLLMAgent: # 保存助手回复到对话 conversation.add_message("assistant", assistant_message) - logger.info("收到回复", length=len(assistant_message)) - return assistant_message + request_id = self._request_id_from_response(response, result) + logger.info("收到回复", length=len(assistant_message), request_id=request_id) + return ModelResult( + content=assistant_message, + usage=self._normalize_usage(result.get("usage")), + request_id=request_id, + response_id=result.get("id"), + model=result.get("model") or self.llm_config.model, + api_format="openai_chat", + endpoint=self.llm_config.chat_endpoint, + ) except httpx.HTTPStatusError as e: logger.error("HTTP错误", status_code=e.response.status_code, detail=e.response.text) @@ -201,6 +324,215 @@ class LiteLLMAgent: except Exception as e: logger.error("请求失败", error=str(e)) raise + + async def _openai_chat_stream_result(self, conversation: Conversation) -> ModelResult: + """OpenAI chat completions stream=true, aggregated into one Runtime artifact.""" + client = await self._get_client() + request_body = { + "model": self.llm_config.model, + "messages": conversation.to_openai_format(), + "temperature": self.llm_config.temperature, + "max_tokens": self.llm_config.max_tokens, + "stream": True, + "stream_options": {"include_usage": True}, + } + + full_response = "" + usage: Dict[str, int] = {} + response_id: Optional[str] = None + request_id: Optional[str] = None + try: + async with client.stream("POST", self.llm_config.chat_endpoint, json=request_body) as response: + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + self._raise_gateway_error(exc) + request_id = self._request_id_from_response(response) + async for line in response.aiter_lines(): + if not line.startswith("data: "): + continue + data = line[6:] + if data == "[DONE]": + break + try: + chunk = json.loads(data) + except json.JSONDecodeError: + continue + response_id = response_id or chunk.get("id") + request_id = request_id or chunk.get("request_id") + if chunk.get("usage"): + usage = self._normalize_usage(chunk.get("usage")) + choices = chunk.get("choices") or [] + if not choices: + continue + delta = choices[0].get("delta", {}) + content = delta.get("content", "") + if content: + full_response += content + + conversation.add_message("assistant", full_response) + logger.info("收到流式回复", length=len(full_response), request_id=request_id) + return ModelResult( + content=full_response, + usage=usage, + request_id=request_id or response_id, + response_id=response_id, + model=self.llm_config.model, + api_format="openai_chat", + endpoint=self.llm_config.chat_endpoint, + ) + except Exception as e: + logger.error("流式请求失败", error=str(e)) + raise + + def _anthropic_payload(self, conversation: Conversation, *, stream: bool = False) -> Dict[str, Any]: + system_parts: List[str] = [] + messages: List[Dict[str, str]] = [] + for message in conversation.messages: + if message.role == "system": + system_parts.append(message.content) + else: + role = "assistant" if message.role == "assistant" else "user" + messages.append({"role": role, "content": message.content}) + payload: Dict[str, Any] = { + "model": self.llm_config.model, + "messages": messages, + "max_tokens": self.llm_config.max_tokens, + "stream": stream, + } + if system_parts: + payload["system"] = "\n\n".join(system_parts) + return payload + + def _anthropic_text(self, body: Dict[str, Any]) -> str: + content = body.get("content") or [] + texts = [ + part.get("text", "") + for part in content + if isinstance(part, dict) and part.get("type") == "text" + ] + return "".join(texts) + + async def _anthropic_messages(self, conversation: Conversation) -> ModelResult: + """Anthropic Messages-compatible call for Claude models.""" + client = await self._get_client() + try: + response = await client.post(self.llm_config.messages_endpoint, json=self._anthropic_payload(conversation)) + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + self._raise_gateway_error(exc) + result = response.json() + assistant_message = self._anthropic_text(result) + conversation.add_message("assistant", assistant_message) + request_id = self._request_id_from_response(response, result) + logger.info("收到 Claude Messages 回复", length=len(assistant_message), request_id=request_id) + return ModelResult( + content=assistant_message, + usage=self._normalize_usage(result.get("usage")), + request_id=request_id, + response_id=result.get("id"), + model=result.get("model") or self.llm_config.model, + api_format="anthropic_messages", + endpoint=self.llm_config.messages_endpoint, + ) + except Exception as e: + logger.error("Claude Messages 请求失败", error=str(e)) + raise + + async def _anthropic_messages_stream(self, conversation: Conversation) -> ModelResult: + """Anthropic Messages stream=true, aggregated into one Runtime artifact.""" + client = await self._get_client() + full_response = "" + usage: Dict[str, int] = {} + response_id: Optional[str] = None + request_id: Optional[str] = None + try: + async with client.stream( + "POST", + self.llm_config.messages_endpoint, + json=self._anthropic_payload(conversation, stream=True), + ) as response: + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + self._raise_gateway_error(exc) + request_id = self._request_id_from_response(response) + async for line in response.aiter_lines(): + if not line.startswith("data: "): + continue + data = line[6:] + if data == "[DONE]": + break + try: + event = json.loads(data) + except json.JSONDecodeError: + continue + event_type = event.get("type") + if event_type == "message_start": + message = event.get("message") or {} + response_id = response_id or message.get("id") + usage = self._normalize_usage(message.get("usage")) + elif event_type == "content_block_delta": + delta = event.get("delta") or {} + text = delta.get("text", "") + if text: + full_response += text + elif event_type == "message_delta": + delta_usage = (event.get("usage") or {}) + if delta_usage: + usage = self._normalize_usage({**usage, **delta_usage}) + + conversation.add_message("assistant", full_response) + logger.info("收到 Claude Messages 流式回复", length=len(full_response), request_id=request_id) + return ModelResult( + content=full_response, + usage=usage, + request_id=request_id or response_id, + response_id=response_id, + model=self.llm_config.model, + api_format="anthropic_messages", + endpoint=self.llm_config.messages_endpoint, + ) + except Exception as e: + logger.error("Claude Messages 流式请求失败", error=str(e)) + raise + + async def _stream_anthropic_messages_text(self, conversation: Conversation) -> AsyncGenerator[str, None]: + """Yield text deltas from Anthropic Messages stream for A2A stream clients.""" + client = await self._get_client() + full_response = "" + try: + async with client.stream( + "POST", + self.llm_config.messages_endpoint, + json=self._anthropic_payload(conversation, stream=True), + ) as response: + try: + response.raise_for_status() + except httpx.HTTPStatusError as exc: + self._raise_gateway_error(exc) + async for line in response.aiter_lines(): + if not line.startswith("data: "): + continue + data = line[6:] + if data == "[DONE]": + break + try: + event = json.loads(data) + except json.JSONDecodeError: + continue + if event.get("type") != "content_block_delta": + continue + delta = event.get("delta") or {} + text = delta.get("text", "") + if text: + full_response += text + yield text + conversation.add_message("assistant", full_response) + except Exception as e: + logger.error("Claude Messages 文本流失败", error=str(e)) + raise async def _stream_chat(self, conversation: Conversation) -> AsyncGenerator[str, None]: """流式对话""" @@ -211,7 +543,8 @@ class LiteLLMAgent: "messages": conversation.to_openai_format(), "temperature": self.llm_config.temperature, "max_tokens": self.llm_config.max_tokens, - "stream": True + "stream": True, + "stream_options": {"include_usage": True}, } full_response = "" @@ -232,7 +565,10 @@ class LiteLLMAgent: try: chunk = json.loads(data) - delta = chunk.get("choices", [{}])[0].get("delta", {}) + choices = chunk.get("choices") or [] + if not choices: + continue + delta = choices[0].get("delta", {}) content = delta.get("content", "") if content: full_response += content diff --git a/agent_templates/agents/a2a_litellm_agent/config.py b/agent_templates/agents/a2a_litellm_agent/config.py index 2ec7b98..3a9f46e 100644 --- a/agent_templates/agents/a2a_litellm_agent/config.py +++ b/agent_templates/agents/a2a_litellm_agent/config.py @@ -18,8 +18,9 @@ class LiteLLMConfig: # 基础URL - 用户提供的LiteLLM服务地址 base_url: str = "https://litellm.graystone-fb459c5d.southeastasia.azurecontainerapps.io" - # 完整的chat completions端点 + # 完整的模型端点 chat_endpoint: str = field(init=False) + messages_endpoint: str = field(init=False) # API密钥 - 优先使用传入的，否则从环境变量获取 api_key: Optional[str] = None @@ -38,6 +39,12 @@ class LiteLLMConfig: # 最大token数 max_tokens: int = 4096 + + # API格式：openai_chat 或 anthropic_messages + api_format: str = "openai_chat" + + # 是否强制使用流式请求聚合完整响应 + use_stream: bool = False def __post_init__(self): self.base_url = ( @@ -52,7 +59,35 @@ class LiteLLMConfig: self.api_key = os.getenv("LITELLM_API_KEY") if self.model is None: self.model = os.getenv("MODEL_NAME") or os.getenv("LITELLM_MODEL", "gpt-4") + model_name = (self.model or "").lower() + + self.api_format = ( + os.getenv("LLM_API_FORMAT") + or os.getenv("MODEL_API_FORMAT") + or ("anthropic_messages" if "claude" in model_name else "openai_chat") + ).lower() + + if "gpt-5.4" in model_name: + self.timeout = 600 + self.use_stream = True + if "claude" in model_name: + self.timeout = 600 + self.use_stream = True + + if os.getenv("LITELLM_TIMEOUT") or os.getenv("LLM_TIMEOUT"): + self.timeout = int(os.getenv("LITELLM_TIMEOUT") or os.getenv("LLM_TIMEOUT")) + if os.getenv("LITELLM_MAX_TOKENS") or os.getenv("LLM_MAX_TOKENS"): + self.max_tokens = int(os.getenv("LITELLM_MAX_TOKENS") or os.getenv("LLM_MAX_TOKENS")) + if os.getenv("LITELLM_STREAM") or os.getenv("LLM_STREAM"): + self.use_stream = (os.getenv("LITELLM_STREAM") or os.getenv("LLM_STREAM", "")).lower() in { + "1", + "true", + "yes", + "on", + } + self.chat_endpoint = f"{self.base_url}/chat/completions" + self.messages_endpoint = f"{self.base_url}/messages" def validate(self) -> bool: """验证配置是否完整""" diff --git a/api/swarm/orchestrator.py b/api/swarm/orchestrator.py index f0be9e8..3cd8949 100644 --- a/api/swarm/orchestrator.py +++ b/api/swarm/orchestrator.py @@ -154,7 +154,10 @@ class SwarmOrchestrator: "model": agent.model, "capabilities": agent.capabilities or [], "system_prompt": agent.system_prompt, - "billing_context": (self.swarm.project_context or {}).get("billing_context") or {}, + "billing_context": { + **((self.swarm.project_context or {}).get("billing_context") or {}), + "max_tokens": ((self.swarm.project_context or {}).get("budget") or {}).get("max_tokens"), + }, } return self.k8s_manager.deploy_swarm_agent( self.swarm_id, @@ -211,6 +214,7 @@ class SwarmOrchestrator: await self._emit_artifacts(failure_artifacts) await self._emit_phase("failed", "Swarm execution failed") await self._emit_status("failed", {"error": str(e)}) + await self._emit_usage() raise async def _execute_sequential(self) -> Dict[str, Any]: @@ -459,8 +463,29 @@ class SwarmOrchestrator: if isinstance(response, dict) and response.get("error"): error = response.get("error") or {} message_text = error.get("message") if isinstance(error, dict) else str(error) - raise RuntimeError(message_text or "A2A agent returned an error response") + error_data = error.get("data") if isinstance(error, dict) else {} + if isinstance(error_data, dict): + error_message = SwarmMessage( + message_id=str(uuid.uuid4()), + swarm_id=self.swarm_id, + from_agent_id=client.agent_id, + to_agent_id=None, + message_type="error", + content=message_text or "A2A agent returned an error response", + message_metadata={ + "newapi_request_id": error_data.get("request_id"), + "model_status_code": error_data.get("status_code"), + "model_response_preview": (error_data.get("response_text") or "")[:1000], + }, + ) + self.db.add(error_message) + self.db.commit() + runtime_error = RuntimeError(message_text or "A2A agent returned an error response") + if isinstance(error_data, dict): + setattr(runtime_error, "request_id", error_data.get("request_id")) + raise runtime_error usage = self._extract_usage(response) + model_metadata = self._extract_model_metadata(response) if self.swarm and usage["total_tokens"]: self.swarm.tokens_used += usage["total_tokens"] self.db.commit() @@ -473,6 +498,8 @@ class SwarmOrchestrator: "summary": "Agent task completed", "result_preview": str(response)[:500], "model_usage": usage, + "newapi_request_id": model_metadata.get("newapi_request_id"), + "model_api_format": model_metadata.get("api_format"), }, ) @@ -484,7 +511,10 @@ class SwarmOrchestrator: to_agent_id=None, # To orchestrator message_type="response", content=str(response), - message_metadata={} + message_metadata={ + "model_usage": usage, + **model_metadata, + }, ) self.db.add(response_message) self.db.commit() @@ -512,6 +542,8 @@ class SwarmOrchestrator: "status": "completed", "summary": self._response_summary(response), "runtime_deployment_id": self.swarm_id, + "newapi_request_id": model_metadata.get("newapi_request_id"), + "model_usage": usage, }, ) @@ -532,6 +564,7 @@ class SwarmOrchestrator: "tool_name": "agent_task", "summary": "Agent task failed", "error": str(e), + "newapi_request_id": self._request_id_from_error(e), }, ) await self._emit( @@ -543,6 +576,7 @@ class SwarmOrchestrator: "status": "failed", "summary": str(e), "runtime_deployment_id": self.swarm_id, + "newapi_request_id": self._request_id_from_error(e), }, ) raise @@ -676,13 +710,20 @@ class SwarmOrchestrator: context = self._callback_context() budget = context.get("budget") or {} billing_context = context.get("billing_context") or {} + usage = self._aggregate_model_usage_from_messages() + if self.swarm and usage["total_tokens"] > self.swarm.tokens_used: + self.swarm.tokens_used = usage["total_tokens"] + self.db.commit() + model_tokens = self.swarm.tokens_used if self.swarm else usage["total_tokens"] + if not model_tokens: + return await self._emit( "budget.alert", payload={ "model_id": billing_context.get("default_model_id") or "unknown", - "model_tokens": self.swarm.tokens_used if self.swarm else 0, - "prompt_tokens": 0, - "completion_tokens": 0, + "model_tokens": model_tokens, + "prompt_tokens": usage["prompt_tokens"], + "completion_tokens": usage["completion_tokens"], "model_cost_usd": 0, "runtime_seconds": 0, "cpu_core_seconds": 0, @@ -701,6 +742,27 @@ class SwarmOrchestrator: }, ) + def _aggregate_model_usage_from_messages(self) -> Dict[str, int]: + """Aggregate persisted model usage metadata for Runtime usage callbacks.""" + totals = {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0} + messages = ( + self.db.query(SwarmMessage) + .filter(SwarmMessage.swarm_id == self.swarm_id) + .all() + ) + for message in messages: + usage = (message.message_metadata or {}).get("model_usage") or {} + if not isinstance(usage, dict): + continue + totals["prompt_tokens"] += int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0) + totals["completion_tokens"] += int(usage.get("completion_tokens") or usage.get("output_tokens") or 0) + totals["total_tokens"] += int( + usage.get("total_tokens") + or (usage.get("prompt_tokens") or usage.get("input_tokens") or 0) + + (usage.get("completion_tokens") or usage.get("output_tokens") or 0) + ) + return totals + async def _emit( self, event_type: str, @@ -845,6 +907,42 @@ class SwarmOrchestrator: "total_tokens": total_tokens, } + def _extract_model_metadata(self, response: Any) -> Dict[str, Any]: + """Extract model request metadata from nested A2A responses.""" + metadata = self._find_model_metadata(response) or {} + request_id = ( + metadata.get("newapi_request_id") + or metadata.get("request_id") + or metadata.get("response_id") + ) + return { + "newapi_request_id": request_id, + "response_id": metadata.get("response_id"), + "model": metadata.get("model"), + "api_format": metadata.get("api_format"), + "endpoint": metadata.get("endpoint"), + } + + def _find_model_metadata(self, value: Any) -> Optional[Dict[str, Any]]: + """Find nested model metadata emitted by an A2A agent.""" + if isinstance(value, dict): + if any(key in value for key in ("newapi_request_id", "request_id", "response_id", "api_format")): + return value + for nested in value.values(): + found = self._find_model_metadata(nested) + if found: + return found + elif isinstance(value, list): + for item in value: + found = self._find_model_metadata(item) + if found: + return found + return None + + def _request_id_from_error(self, error: Exception) -> Optional[str]: + """Best-effort extraction for errors raised after model calls.""" + return getattr(error, "request_id", None) + def _find_usage_dict(self, value: Any) -> Optional[Dict[str, Any]]: """Find a nested usage-like dict containing token counters.""" if isinstance(value, dict): diff --git a/api/swarm/router.py b/api/swarm/router.py index 40c0cc7..19fb88f 100644 --- a/api/swarm/router.py +++ b/api/swarm/router.py @@ -416,10 +416,19 @@ async def get_swarm_logs(swarm_id: str, db: Session = Depends(get_db)): if agent.output: log_lines.append(f"last_output={agent.output[:500]}") if agent_messages: - log_lines.extend( - f"{message.created_at.isoformat()} {message.message_type}: {message.content[:300]}" - for message in agent_messages[-5:] - ) + for message in agent_messages[-5:]: + metadata = message.message_metadata or {} + request_id = metadata.get("newapi_request_id") + usage = metadata.get("model_usage") or {} + suffix_parts = [] + if request_id: + suffix_parts.append(f"newapi_request_id={request_id}") + if usage.get("total_tokens"): + suffix_parts.append(f"tokens={usage['total_tokens']}") + suffix = f" ({', '.join(suffix_parts)})" if suffix_parts else "" + log_lines.append( + f"{message.created_at.isoformat()} {message.message_type}: {message.content[:300]}{suffix}" + ) else: log_lines.append("no_runtime_messages_recorded") @@ -429,6 +438,17 @@ async def get_swarm_logs(swarm_id: str, db: Session = Depends(get_db)): "namespace": agent.namespace, "pod_name": agent.pod_name, "logs": "\n".join(log_lines), + "messages": [ + { + "message_id": message.message_id, + "message_type": message.message_type, + "created_at": message.created_at, + "newapi_request_id": (message.message_metadata or {}).get("newapi_request_id"), + "model_usage": (message.message_metadata or {}).get("model_usage"), + "metadata": message.message_metadata or {}, + } + for message in agent_messages[-20:] + ], }) return logs diff --git a/k8s_manager.py b/k8s_manager.py index ce9c80f..36176e0 100644 --- a/k8s_manager.py +++ b/k8s_manager.py @@ -2449,11 +2449,27 @@ echo "Identity volume initialized successfully" role = agent_config.get("role", "worker") model = agent_config.get("model", "gpt-4") billing_context = agent_config.get("billing_context") or {} + model_lower = model.lower() + api_format = ( + billing_context.get("api_format") + or ("anthropic_messages" if "claude" in model_lower else "openai_chat") + ) + stream_enabled = billing_context.get("stream") + if stream_enabled is None: + stream_enabled = "gpt-5.4" in model_lower or "claude" in model_lower + elif isinstance(stream_enabled, str): + stream_enabled = stream_enabled.lower() in {"1", "true", "yes", "on"} + timeout_seconds = int( + billing_context.get("timeout_sec") + or billing_context.get("timeout_seconds") + or (600 if stream_enabled else 300) + ) + max_tokens = int(billing_context.get("max_tokens") or 4096) # Get template image # TODO: Load from template database image_map = { - "a2a_litellm_agent": "agnettaiji.azurecr.io/ai-agents/a2a-litellm-agent:heicode-v2-gpt54-202605300145-arm64", + "a2a_litellm_agent": "agnettaiji.azurecr.io/ai-agents/a2a-litellm-agent:heicode-v2-runtime-20260531214518-arm64", "code_manager_agent": "agnettaiji.azurecr.io/ai-agents/code-manager-agent:latest" } image = image_map.get(template, image_map["a2a_litellm_agent"]) @@ -2466,6 +2482,14 @@ echo "Identity volume initialized successfully" client.V1EnvVar(name="MODEL_NAME", value=model), client.V1EnvVar(name="POD_NAME", value=pod_name), client.V1EnvVar(name="NAMESPACE", value=namespace), + client.V1EnvVar(name="MODEL_API_FORMAT", value=api_format), + client.V1EnvVar(name="LLM_API_FORMAT", value=api_format), + client.V1EnvVar(name="LITELLM_STREAM", value="true" if stream_enabled else "false"), + client.V1EnvVar(name="LLM_STREAM", value="true" if stream_enabled else "false"), + client.V1EnvVar(name="LITELLM_TIMEOUT", value=str(timeout_seconds)), + client.V1EnvVar(name="LLM_TIMEOUT", value=str(timeout_seconds)), + client.V1EnvVar(name="LITELLM_MAX_TOKENS", value=str(max_tokens)), + client.V1EnvVar(name="LLM_MAX_TOKENS", value=str(max_tokens)), ] gateway_url = ( diff --git a/scripts/e2e_complete_project.py b/scripts/e2e_complete_project.py new file mode 100755 index 0000000..4ae71f3 --- /dev/null +++ b/scripts/e2e_complete_project.py @@ -0,0 +1,311 @@ +#!/usr/bin/env python3 +"""End-to-end sub-mode Runtime check for a complete generated project. + +This script is intentionally outside the normal unit-test suite because it +requires a live agent-manager Runtime, Kubernetes, and a model gateway. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import re +import shutil +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from pathlib import Path +from typing import Any + + +DEFAULT_OBJECTIVE = """Return ONLY valid JSON: +{"files":[{"path":"...","content":"..."}],"run_tests":"python -m unittest discover -s tests -v","smoke_test":"python -m textstats_cli samples/example.txt --json"}. + +Build a complete, compact Python stdlib project named textstats_cli. +Include exactly these files: +- pyproject.toml +- README.md +- textstats_cli/__main__.py +- textstats_cli/core.py +- tests/test_core.py +- samples/example.txt + +Features: +- CLI accepts a text file path. +- --json outputs JSON. +- Default output is human readable. +- Report line count, word count, character count, top 5 words excluding common stopwords, and estimated reading time. +- Tests must cover counting, stopword filtering, JSON-safe result shape, and missing-file error handling. + +Constraints: +- Use only the Python standard library. +- Keep the project small enough for one response. +- No placeholders. +- No markdown fences. +- No prose outside the JSON object. +""" + + +TERMINAL_STATUSES = {"completed", "failed", "stopped"} + + +def http_json(method: str, url: str, payload: dict[str, Any] | None = None, headers: dict[str, str] | None = None) -> Any: + data = None if payload is None else json.dumps(payload).encode("utf-8") + request_headers = {"Content-Type": "application/json", **(headers or {})} + request = urllib.request.Request(url, data=data, headers=request_headers, method=method) + try: + with urllib.request.urlopen(request, timeout=60) as response: + return json.loads(response.read().decode("utf-8")) + except urllib.error.HTTPError as exc: + body = exc.read().decode("utf-8", errors="replace") + raise RuntimeError(f"{method} {url} failed with HTTP {exc.code}: {body}") from exc + + +def http_bytes(url: str) -> bytes: + with urllib.request.urlopen(url, timeout=120) as response: + return response.read() + + +def extract_json_object(text: str) -> dict[str, Any]: + raw = text.strip() + fenced = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", raw, re.S) + if fenced: + raw = fenced.group(1) + else: + start = raw.find("{") + end = raw.rfind("}") + if start < 0 or end <= start: + raise ValueError("artifact does not contain a JSON object") + raw = raw[start : end + 1] + return json.loads(raw) + + +def safe_write_project(project_dir: Path, files: list[dict[str, Any]]) -> None: + project_root = project_dir.resolve() + for item in files: + relative_path = item.get("path") + content = item.get("content") + if not isinstance(relative_path, str) or not relative_path: + raise ValueError(f"invalid file path in artifact: {item!r}") + if not isinstance(content, str): + raise ValueError(f"invalid content for {relative_path}") + target = (project_root / relative_path).resolve() + if project_root not in target.parents and target != project_root: + raise ValueError(f"artifact attempted to write outside project: {relative_path}") + target.parent.mkdir(parents=True, exist_ok=True) + target.write_text(content, encoding="utf-8") + + +def run_command(command: str, cwd: Path) -> subprocess.CompletedProcess[str]: + env = os.environ.copy() + shim_dir = None + if shutil.which("python") is None: + shim_dir = Path(tempfile.mkdtemp(prefix="heicode-python-shim-")) + (shim_dir / "python").symlink_to(sys.executable) + env["PATH"] = str(shim_dir) + os.pathsep + env.get("PATH", "") + return subprocess.run( + command, + cwd=str(cwd), + env=env, + shell=True, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + timeout=120, + check=False, + ) + + +def build_payload(args: argparse.Namespace) -> dict[str, Any]: + return { + "orchestration_plan": { + "sub_mode": "agile", + "objective": args.objective, + "user_context": {"user_id": args.user_id}, + "agents": [ + { + "role": "backend", + "template": "a2a_litellm_agent", + "model": args.model, + "capabilities": ["code", "test"], + } + ], + "agile_context": {"max_iterations": 1, "stage": "development"}, + "budget": {"max_duration_sec": args.timeout_seconds, "max_tokens": args.max_tokens}, + "billing_context": { + "provider": "newapi", + "default_model_id": args.model, + "model_gateway_url": args.model_gateway_url, + "api_format": args.api_format, + "stream": args.stream, + "timeout_sec": args.model_timeout_seconds, + "max_tokens": args.max_tokens, + }, + "metadata": {"correlation_id": args.correlation_id}, + }, + "callback": {"url": args.callback_url, "method": "POST"}, + } + + +def poll_swarm(base_url: str, swarm_id: str, timeout_seconds: int, poll_interval: int) -> dict[str, Any]: + deadline = time.time() + timeout_seconds + last_line = None + while time.time() < deadline: + status = http_json("GET", f"{base_url}/api/swarms/{swarm_id}") + line = { + "status": status.get("status"), + "phase": status.get("phase"), + "progress": status.get("progress"), + "artifact_count": len(status.get("artifacts") or []), + "tokens_used": (status.get("metrics") or {}).get("tokens_used"), + } + if line != last_line: + print("status:", json.dumps(line, ensure_ascii=False)) + last_line = line + if status.get("status") in TERMINAL_STATUSES: + return status + time.sleep(poll_interval) + raise TimeoutError(f"swarm {swarm_id} did not finish within {timeout_seconds}s") + + +def assert_runtime_observability(base_url: str, swarm_id: str, status: dict[str, Any]) -> None: + metrics = status.get("metrics") or {} + tokens_used = int(metrics.get("tokens_used") or 0) + if tokens_used <= 0: + raise AssertionError(f"Runtime tokens_used must be > 0, got {tokens_used}") + + logs = http_json("GET", f"{base_url}/api/swarms/{swarm_id}/logs") + request_ids: list[str] = [] + usage_totals: list[int] = [] + for agent in logs.get("agents") or []: + for message in agent.get("messages") or []: + if message.get("newapi_request_id"): + request_ids.append(message["newapi_request_id"]) + usage = message.get("model_usage") or {} + if usage.get("total_tokens"): + usage_totals.append(int(usage["total_tokens"])) + if not request_ids: + raise AssertionError("Runtime logs must include at least one NewAPI request_id") + if not usage_totals: + raise AssertionError("Runtime logs must include model usage with total_tokens") + print("observability:", json.dumps({"request_ids": request_ids, "usage_totals": usage_totals}, ensure_ascii=False)) + + +def run_project_validation(base_url: str, status: dict[str, Any], output_dir: Path) -> None: + artifacts = status.get("artifacts") or [] + if status.get("status") != "completed": + raise AssertionError(f"swarm did not complete: {status.get('error_message')}") + if not artifacts: + raise AssertionError("completed swarm returned no artifacts") + + artifact = artifacts[0] + download_path = (artifact.get("metadata") or {}).get("download_path") + if not download_path: + raise AssertionError("artifact metadata is missing download_path") + + artifact_bytes = http_bytes(f"{base_url}{download_path}") + artifact_hash = "sha256:" + hashlib.sha256(artifact_bytes).hexdigest() + expected_hash = (artifact.get("metadata") or {}).get("content_hash") + if expected_hash and artifact_hash != expected_hash: + raise AssertionError(f"artifact hash mismatch: expected {expected_hash}, got {artifact_hash}") + + artifact_text = artifact_bytes.decode("utf-8") + artifact_json = extract_json_object(artifact_text) + files = artifact_json.get("files") + if not isinstance(files, list) or len(files) < 5: + raise AssertionError("artifact must contain a multi-file project") + + project_dir = output_dir / "project" + if project_dir.exists(): + shutil.rmtree(project_dir) + project_dir.mkdir(parents=True) + safe_write_project(project_dir, files) + + required_paths = { + "pyproject.toml", + "README.md", + "textstats_cli/__main__.py", + "textstats_cli/core.py", + "tests/test_core.py", + "samples/example.txt", + } + actual_paths = {str(path.relative_to(project_dir)) for path in project_dir.rglob("*") if path.is_file()} + missing = sorted(required_paths - actual_paths) + if missing: + raise AssertionError(f"generated project missing required files: {missing}") + + for label, command in ( + ("run_tests", artifact_json.get("run_tests")), + ("smoke_test", artifact_json.get("smoke_test")), + ): + if not isinstance(command, str) or not command.strip(): + raise AssertionError(f"artifact missing {label}") + result = run_command(command, project_dir) + print(f"{label}: {command}") + print(result.stdout) + if result.returncode != 0: + raise AssertionError(f"{label} failed with exit code {result.returncode}") + + print("artifact:", json.dumps({ + "artifact_id": artifact.get("artifact_id"), + "uri": artifact.get("uri"), + "size_bytes": artifact.get("size_bytes"), + "content_hash": artifact_hash, + "project_dir": str(project_dir), + }, ensure_ascii=False)) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--base-url", default=os.getenv("AGENT_MANAGER_URL", "http://127.0.0.1:8000")) + parser.add_argument("--model-gateway-url", default=os.getenv("HEICODE_NEWAPI_BASE_URL", "https://code.xinghanlab.com/v1")) + parser.add_argument("--model", default=os.getenv("HEICODE_E2E_MODEL", "gpt-5.4")) + parser.add_argument("--api-format", default=os.getenv("HEICODE_E2E_API_FORMAT", "openai_chat")) + parser.add_argument("--stream", action=argparse.BooleanOptionalAction, default=True) + parser.add_argument("--max-tokens", type=int, default=int(os.getenv("HEICODE_E2E_MAX_TOKENS", "6000"))) + parser.add_argument("--timeout-seconds", type=int, default=int(os.getenv("HEICODE_E2E_TIMEOUT_SECONDS", "1200"))) + parser.add_argument("--model-timeout-seconds", type=int, default=int(os.getenv("HEICODE_E2E_MODEL_TIMEOUT_SECONDS", "600"))) + parser.add_argument("--poll-interval", type=int, default=10) + parser.add_argument("--user-id", default="heicode-complete-project-e2e") + parser.add_argument("--callback-url", default="http://127.0.0.1:9/heicode-callback") + parser.add_argument("--correlation-id", default=f"heicode-complete-project-e2e-{int(time.time())}") + parser.add_argument("--idempotency-key", default=f"complete-project-e2e-{int(time.time())}") + parser.add_argument("--objective", default=DEFAULT_OBJECTIVE) + parser.add_argument("--output-dir", type=Path, default=None) + return parser.parse_args() + + +def main() -> int: + args = parse_args() + base_url = args.base_url.rstrip("/") + output_dir = args.output_dir or Path(tempfile.mkdtemp(prefix="heicode-complete-project-e2e-")) + output_dir.mkdir(parents=True, exist_ok=True) + + payload = build_payload(args) + created = http_json( + "POST", + f"{base_url}/api/swarms", + payload, + headers={"X-Idempotency-Key": args.idempotency_key}, + ) + swarm_id = created["swarm_id"] + print("created:", json.dumps({"swarm_id": swarm_id, "output_dir": str(output_dir)}, ensure_ascii=False)) + + status = poll_swarm(base_url, swarm_id, args.timeout_seconds, args.poll_interval) + run_project_validation(base_url, status, output_dir) + assert_runtime_observability(base_url, swarm_id, status) + print("PASS: complete project E2E succeeded") + return 0 + + +if __name__ == "__main__": + try: + raise SystemExit(main()) + except Exception as exc: + print(f"FAIL: {exc}", file=sys.stderr) + raise SystemExit(1)