diff --git a/docs/design/cloud-agents/e2e-testing-strategy.md b/docs/design/cloud-agents/e2e-testing-strategy.md index 50b0e2879..88eae6c5f 100644 --- a/docs/design/cloud-agents/e2e-testing-strategy.md +++ b/docs/design/cloud-agents/e2e-testing-strategy.md @@ -159,10 +159,8 @@ name in two different files). Current layout: | File | Layer | Covers | |---|---|---| | `test_workflows_http_e2e.py` | real HTTP (`TestClient`) | `/v1/workflows/*`, spawn none+local+ephemeral, incl. one-step workflows | -| `test_query_direct_handler_e2e.py` | handler-direct | `/v1/query/direct` error paths | | `test_step_executor_e2e.py` | step-executor dispatch (`get_step_executor(...).run(...)`) | single-step execution, spawn none+local+ephemeral | | `test_workflow_definitions_e2e.py` | step-executor dispatch | full workflow-YAML execution, one step-executor call per step | -| `test_otel_tracing_e2e.py` | mid-layer (`execute_query_via_direct_executor`) | trace/span assertions for the query/direct path | | `test_workflow_tracing_e2e.py` | `LocalWorkflowRunner` directly | trace/span assertions for the workflow-engine path | | `mock_llm_server.py`, `mock_llm_env.py`, `test_mock_llm_*.py`, `jaeger_helpers.py`, `conftest.py` | infra | shared fixtures/mocks, not endpoint tests themselves | diff --git a/docs/design/cloud-agents/integration-architecture.md b/docs/design/cloud-agents/integration-architecture.md index 110165986..d9eb2f02f 100644 --- a/docs/design/cloud-agents/integration-architecture.md +++ b/docs/design/cloud-agents/integration-architecture.md @@ -14,14 +14,12 @@ graph TB a2a["/a2a"] end subgraph new["New Endpoints"] - qd["/query/direct"] - qds["/query/direct/stream"] wf["/workflows/*"] at["/agent-tools"] end existing -->|build_agent| llamastack["Llama Stack"] - new --> bridge["query_executor.py\nvalidate + build StepInput"] + new --> bridge["workflow/\nstorage + executor_factory"] end bridge --> dispatch["get_step_executor()"] @@ -222,7 +220,6 @@ graph TB end subgraph bridge["Bridge Layer (lightspeed-stack src/workflow/)"] - qe["query_executor.py"] storage["storage.py"] ef["executor_factory.py"] end @@ -264,8 +261,6 @@ graph TB | Endpoint | Method | Purpose | |---|---|---| -| `/v1/query/direct` | POST | Blocking query via DirectExecutor | -| `/v1/query/direct/stream` | POST | SSE streaming query | | `/v1/workflows/run` | POST | Start a workflow (multi-step, or one-step for one-shot agent runs) | | `/v1/workflows/{id}` | GET | Get workflow status | | `/v1/workflows/{id}/approve` | POST | Approve a paused step | @@ -375,6 +370,10 @@ graph LR style future fill:#e8f5e9,stroke:#4caf50 ``` +> **Deferred:** the `/v1/query/direct` endpoints were removed to keep the +> workflow work focused; bring-back is tracked in jameswnl/lightspeed-stack#59. +> The status below describes the design target, not shipped endpoints. + ### What's Done | Component | Status | @@ -475,7 +474,6 @@ Blue = cloud-agents. Orange = lightspeed-stack. Red = final migration steps. | Component | Multi-pod safe? | Notes | |---|---|---| -| `/query/direct` | Yes | Stateless per call | | Conversation state | Yes (PostgreSQL) | Shared database | | Workflow state | Yes (PostgreSQL) | Shared database | | Running workflow tasks | No (in-memory) | Use Temporal for crash recovery | diff --git a/docs/devel_doc/openapi.json b/docs/devel_doc/openapi.json index 813f0c7e0..8737dbf3f 100644 --- a/docs/devel_doc/openapi.json +++ b/docs/devel_doc/openapi.json @@ -6239,90 +6239,6 @@ } } }, - "/v1/query/direct": { - "post": { - "tags": [ - "query" - ], - "summary": "Query Direct Handler", - "description": "Execute a query via DirectExecutor (no Llama Stack).\n\nParameters:\n request: FastAPI request (consumed by @authorize decorator).\n body: Query parameters (field names match /query).\n auth: Authentication tuple (consumed by @authorize decorator).\n\nReturns:\n Response matching QueryResponse field names.", - "operationId": "query_direct_handler_v1_query_direct_post", - "requestBody": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/QueryDirectRequest" - } - } - }, - "required": true - }, - "responses": { - "200": { - "description": "Successful Response", - "content": { - "application/json": { - "schema": { - "additionalProperties": true, - "type": "object", - "title": "Response Query Direct Handler V1 Query Direct Post" - } - } - } - }, - "422": { - "description": "Validation Error", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - } - } - } - } - }, - "/v1/query/direct/stream": { - "post": { - "tags": [ - "query" - ], - "summary": "Query Direct Stream Handler", - "description": "Stream a query via DirectExecutor as SSE events.\n\nToken-by-token streaming is active when tools or MCP servers\nare configured. Without tools, yields a single complete event.\n\nParameters:\n request: FastAPI request (consumed by @authorize decorator).\n body: Query parameters (field names match /query).\n auth: Authentication tuple (consumed by @authorize decorator).\n\nReturns:\n StreamingResponse with SSE events.", - "operationId": "query_direct_stream_handler_v1_query_direct_stream_post", - "requestBody": { - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/QueryDirectRequest" - } - } - }, - "required": true - }, - "responses": { - "200": { - "description": "Successful Response", - "content": { - "application/json": { - "schema": {} - } - } - }, - "422": { - "description": "Validation Error", - "content": { - "application/json": { - "schema": { - "$ref": "#/components/schemas/HTTPValidationError" - } - } - } - } - } - } - }, "/v1/streaming_query": { "post": { "tags": [ @@ -18889,100 +18805,6 @@ } ] }, - "QueryDirectRequest": { - "properties": { - "query": { - "type": "string", - "title": "Query", - "description": "The query string", - "examples": [ - "What is Kubernetes?" - ] - }, - "conversation_id": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Conversation Id", - "description": "Conversation ID for multi-turn (requires PostgreSQL transcript store)" - }, - "provider": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Provider", - "description": "The optional provider" - }, - "model": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "Model", - "description": "The optional model" - }, - "system_prompt": { - "anyOf": [ - { - "type": "string" - }, - { - "type": "null" - } - ], - "title": "System Prompt", - "description": "The optional system prompt" - }, - "mcp_servers": { - "anyOf": [ - { - "items": { - "type": "string" - }, - "type": "array" - }, - { - "type": "null" - } - ], - "title": "Mcp Servers", - "description": "MCP server names from config (None = all configured)" - }, - "output_schema": { - "anyOf": [ - { - "additionalProperties": true, - "type": "object" - }, - { - "type": "null" - } - ], - "title": "Output Schema", - "description": "JSON Schema for structured output" - } - }, - "type": "object", - "required": [ - "query" - ], - "title": "QueryDirectRequest", - "description": "Request body for POST /v1/query/direct.\n\nField names match QueryRequest for API compatibility.\n\nAttributes:\n query: The query string.\n conversation_id: Optional conversation ID for multi-turn.\n provider: Optional provider name.\n model: Optional model name.\n system_prompt: Optional system prompt.\n mcp_servers: Optional MCP server names from config.\n output_schema: Optional JSON Schema for structured output." - }, "QueryRequest": { "properties": { "query": { diff --git a/src/app/endpoints/query_direct.py b/src/app/endpoints/query_direct.py deleted file mode 100644 index cd7eb44c0..000000000 --- a/src/app/endpoints/query_direct.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Handler for /query/direct — query via DirectExecutor (no Llama Stack). - -Parallel endpoint to /query that uses cloud-agents' DirectExecutor -instead of build_agent() → Llama Stack. Uses the SAME request field -names as /query so callers can switch by changing the URL. - -Current limitations (tracked in issue #9): -- No shield moderation, RAG, compaction -- conversation_id is NOT wired — multi-turn state is not managed -- Streaming yields token events only when tools/MCP are configured; - plain prompts produce a single complete event -""" - -import json -from collections.abc import AsyncIterator -from typing import Annotated, Any, Optional - -from fastapi import APIRouter, Depends, HTTPException, Request, status -from fastapi.responses import StreamingResponse -from pydantic import BaseModel, Field - -from authentication import get_auth_dependency -from authentication.interface import AuthTuple -from authorization.middleware import authorize -from configuration import configuration -from log import get_logger -from models.config import Action -from utils.endpoints import check_configuration_loaded -from workflow.query_executor import ( - execute_query_via_direct_executor, - stream_query_via_direct_executor, -) - -logger = get_logger(__name__) -router = APIRouter(tags=["query"]) - - -class QueryDirectRequest(BaseModel): - """Request body for POST /v1/query/direct. - - Field names match QueryRequest for API compatibility. - - Attributes: - query: The query string. - conversation_id: Optional conversation ID for multi-turn. - provider: Optional provider name. - model: Optional model name. - system_prompt: Optional system prompt. - mcp_servers: Optional MCP server names from config. - output_schema: Optional JSON Schema for structured output. - """ - - query: str = Field( - ..., - description="The query string", - examples=["What is Kubernetes?"], - ) - - conversation_id: Optional[str] = Field( - None, - description="Conversation ID for multi-turn (requires PostgreSQL transcript store)", - ) - - provider: Optional[str] = Field( - None, - description="The optional provider", - ) - - model: Optional[str] = Field( - None, - description="The optional model", - ) - - system_prompt: Optional[str] = Field( - None, - description="The optional system prompt", - ) - - mcp_servers: Optional[list[str]] = Field( - None, - description="MCP server names from config (None = all configured)", - ) - - output_schema: Optional[dict[str, Any]] = Field( - None, - description="JSON Schema for structured output", - ) - - -def _serialize_output(output: Any) -> str: - """Serialize executor output to a response string. - - Parameters: - output: StepResult.output — can be dict, str, int, or None. - - Returns: - String representation of the output. - """ - if output is None: - return "" - if isinstance(output, dict): - return output.get("response", str(output)) - return str(output) - - -@router.post("/query/direct") -@authorize(Action.QUERY) -async def query_direct_handler( - request: Request, - body: QueryDirectRequest, - auth: Annotated[AuthTuple, Depends(get_auth_dependency())], -) -> dict[str, Any]: - """Execute a query via DirectExecutor (no Llama Stack). - - Parameters: - request: FastAPI request (consumed by @authorize decorator). - body: Query parameters (field names match /query). - auth: Authentication tuple (consumed by @authorize decorator). - - Returns: - Response matching QueryResponse field names. - """ - _ = request - user_id, username, _, _ = auth - - check_configuration_loaded(configuration) - - try: - result = await execute_query_via_direct_executor( - prompt=body.query, - model=body.model, - provider=body.provider, - instructions=body.system_prompt, - mcp_server_names=body.mcp_servers, - output_schema=body.output_schema, - conversation_id=body.conversation_id, - user_id=user_id, - username=username, - ) - except ValueError as exc: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=str(exc), - ) from exc - - return { - "conversation_id": body.conversation_id, - "response": _serialize_output(result.output), - "truncated": False, - "input_tokens": result.input_tokens, - "output_tokens": result.output_tokens, - "available_quotas": {}, - "tool_calls": [], - "tool_results": [], - "rag_chunks": [], - "referenced_documents": [], - "request_id": "", - "interrupted": False, - } - - -@router.post("/query/direct/stream") -@authorize(Action.STREAMING_QUERY) -async def query_direct_stream_handler( - request: Request, - body: QueryDirectRequest, - auth: Annotated[AuthTuple, Depends(get_auth_dependency())], -) -> StreamingResponse: - """Stream a query via DirectExecutor as SSE events. - - Token-by-token streaming is active when tools or MCP servers - are configured. Without tools, yields a single complete event. - - Parameters: - request: FastAPI request (consumed by @authorize decorator). - body: Query parameters (field names match /query). - auth: Authentication tuple (consumed by @authorize decorator). - - Returns: - StreamingResponse with SSE events. - """ - _ = request - user_id, username, _, _ = auth - - check_configuration_loaded(configuration) - - try: - from workflow.query_executor import ( # pylint: disable=import-outside-toplevel - _resolve_provider, - _validate_prompt, - ) - - _validate_prompt(body.query, body.system_prompt) - _resolve_provider(body.provider, body.model) - except ValueError as exc: - raise HTTPException( - status_code=status.HTTP_400_BAD_REQUEST, - detail=str(exc), - ) from exc - - stream_args = dict( - prompt=body.query, - model=body.model, - provider=body.provider, - instructions=body.system_prompt, - mcp_server_names=body.mcp_servers, - output_schema=body.output_schema, - user_id=user_id, - username=username, - ) - - async def event_generator() -> AsyncIterator[str]: - """Generate SSE events from DirectExecutor stream.""" - try: - async for event in stream_query_via_direct_executor( - **stream_args, - ): - event_data: dict[str, Any] = { - "type": event.type, - "data": event.data, - } - if event.result: - event_data["result"] = { - "status": event.result.status, - "output": _serialize_output(event.result.output), - "input_tokens": event.result.input_tokens, - "output_tokens": event.result.output_tokens, - "duration_ms": event.result.duration_ms, - } - yield f"event: {event.type}\ndata: {json.dumps(event_data, default=str)}\n\n" - except Exception as exc: - logger.error("Stream error: %s", exc, exc_info=True) - error_event = { - "type": "error", - "data": {"message": str(exc)}, - } - yield f"event: error\ndata: {json.dumps(error_event)}\n\n" - - return StreamingResponse( - event_generator(), - media_type="text/event-stream", - headers={ - "Cache-Control": "no-cache", - "X-Accel-Buffering": "no", - }, - ) diff --git a/src/app/routers.py b/src/app/routers.py index 9e653012c..da22c6a82 100644 --- a/src/app/routers.py +++ b/src/app/routers.py @@ -21,7 +21,6 @@ providers, # Query endpoints for Response API support query, - query_direct, rags, responses, # RHEL Lightspeed rlsapi v1 compatibility @@ -66,7 +65,6 @@ def include_routers(app: FastAPI) -> None: app.include_router(vector_stores.router, prefix="/v1") # Query endpoints app.include_router(query.router, prefix="/v1") - app.include_router(query_direct.router, prefix="/v1") app.include_router(streaming_query.router, prefix="/v1") app.include_router(stream_interrupt.router, prefix="/v1") app.include_router(config.router, prefix="/v1") diff --git a/src/workflow/middleware.py b/src/workflow/middleware.py deleted file mode 100644 index c67853cb3..000000000 --- a/src/workflow/middleware.py +++ /dev/null @@ -1,129 +0,0 @@ -"""StepMiddleware implementations for lightspeed-stack. - -Cross-cutting concerns that wrap executor calls via cloud-agents' -StepMiddleware protocol. Injected into ChatWorkflowRunner via the -middlewares parameter. -""" - -from __future__ import annotations - -from cloud_agents.workflow.executor.step.base import StepInput, StepResult -from opentelemetry import trace - -from log import get_logger -from utils.otel_tracing import SpanAttributes, anonymize_value - -logger = get_logger(__name__) - - -class TracingMiddleware: - """OTEL tracing middleware — records attributes on the active span. - - Does NOT create spans. The MiddlewareExecutor (cloud-agents) creates - the span; this middleware enriches it with lightspeed-stack-specific - attributes using OTel semantic conventions. - """ - - async def before(self, step_input: StepInput) -> StepInput: - """No-op before hook.""" - return step_input - - async def after(self, step_input: StepInput, result: StepResult) -> StepResult: - """Record step metrics on the current span.""" - span = trace.get_current_span() - if not span.is_recording(): - return result - - provider = step_input.provider or {} - provider_name = provider.get("name", "") - model_name = provider.get("model", "") - if provider_name or model_name: - span.set_attribute( - SpanAttributes.LLM_MODEL_ID, - f"{provider_name}:{model_name}", - ) - span.set_attribute(SpanAttributes.LLM_USAGE_INPUT_TOKENS, result.input_tokens) - span.set_attribute(SpanAttributes.LLM_USAGE_OUTPUT_TOKENS, result.output_tokens) - - if step_input.metadata and step_input.metadata.user_id: - span.set_attribute( - SpanAttributes.USER_ID, - anonymize_value(step_input.metadata.user_id), - ) - - return result - - -class AuditMiddleware: - """Structured audit logging for every agent execution.""" - - async def before(self, step_input: StepInput) -> StepInput: - """Log step start.""" - user = step_input.metadata.user_id if step_input.metadata else "anonymous" - logger.info( - "Agent execution started: step=%s user=%s", - step_input.step_name, - user, - ) - return step_input - - async def after(self, step_input: StepInput, result: StepResult) -> StepResult: - """Log step completion with metrics.""" - user = step_input.metadata.user_id if step_input.metadata else "anonymous" - logger.info( - "Agent execution completed: step=%s user=%s status=%s " - "tokens_in=%d tokens_out=%d duration_ms=%d", - step_input.step_name, - user, - result.status, - result.input_tokens, - result.output_tokens, - result.duration_ms, - ) - return result - - -class QuotaMiddleware: - """Token quota enforcement middleware. - - Placeholder implementation that logs quota checks. Will be wired - to lightspeed-stack's quota infrastructure (src/quota/) when the - quota path is decoupled from Llama Stack. - """ - - async def before(self, step_input: StepInput) -> StepInput: - """Check quota before execution.""" - user = step_input.metadata.user_id if step_input.metadata else None - if user: - logger.debug("Quota check: user=%s (enforcement pending)", user) - return step_input - - async def after(self, step_input: StepInput, result: StepResult) -> StepResult: - """Deduct tokens after execution.""" - user = step_input.metadata.user_id if step_input.metadata else None - # Track consumption regardless of status — failed steps still use tokens - if user: - logger.debug( - "Quota deduct: user=%s status=%s tokens=%d (enforcement pending)", - user, - result.status, - result.input_tokens + result.output_tokens, - ) - return result - - -def get_default_middleware() -> list: - """Return the default middleware stack for lightspeed-stack. - - Order matters: outermost first. TracingMiddleware enriches the - span created by cloud-agents' MiddlewareExecutor. AuditMiddleware - logs before/after. QuotaMiddleware will enforce limits. - - Returns: - List of middleware instances. - """ - return [ - TracingMiddleware(), - AuditMiddleware(), - QuotaMiddleware(), - ] diff --git a/src/workflow/query_executor.py b/src/workflow/query_executor.py deleted file mode 100644 index 8d1aacee8..000000000 --- a/src/workflow/query_executor.py +++ /dev/null @@ -1,309 +0,0 @@ -"""Bridge between /query endpoint and cloud-agents' ChatWorkflowRunner. - -Manages a ChatWorkflowRunner instance that handles multi-turn -conversation state, middleware, and executor dispatch. This module -is the migration path from build_agent() → Llama Stack to -ChatWorkflowRunner → pydantic-ai for the /query chat agent. -""" - -from __future__ import annotations - -from collections.abc import AsyncIterator -from typing import Any, Optional - -from cloud_agents.workflow.executor.chat.runner import ( - ChatWorkflowConfig, - ChatWorkflowRunner, -) -from cloud_agents.workflow.executor.step.base import ( - StepResult, - StreamEvent, -) - -from configuration import configuration -from log import get_logger -from workflow.storage import WorkflowStorageFactory - -logger = get_logger(__name__) - -MAX_PROMPT_LENGTH = 100_000 -MAX_INSTRUCTIONS_LENGTH = 50_000 - -_runner: Optional[ChatWorkflowRunner] = None - - -def resolve_mcp_servers( - server_names: Optional[list[str]] = None, -) -> list[dict[str, Any]]: - """Resolve MCP server names to configs for the executor. - - Parameters: - server_names: Optional list of server names to resolve. - If None, all configured servers are included. - - Returns: - List of MCP server config dicts. - - Raises: - ValueError: If a requested server name is not configured. - """ - configured = {s.name: s for s in configuration.mcp_servers} - - if server_names: - unknown = [n for n in server_names if n not in configured] - if unknown: - raise ValueError( - f"Unknown MCP server(s): {unknown}. " - f"Configured: {sorted(configured.keys())}" - ) - - mcp_configs: list[dict[str, Any]] = [] - for server in configuration.mcp_servers: - if server_names and server.name not in server_names: - continue - config: dict[str, Any] = { - "name": server.name, - "url": server.url, - } - if server.resolved_authorization_headers: - config["headers"] = dict(server.resolved_authorization_headers) - mcp_configs.append(config) - return mcp_configs - - -def _validate_prompt(prompt: str, instructions: Optional[str] = None) -> None: - """Validate prompt and instructions length. - - Raises: - ValueError: On validation failure. - """ - if len(prompt) > MAX_PROMPT_LENGTH: - raise ValueError( - f"Prompt exceeds maximum length ({len(prompt)} > {MAX_PROMPT_LENGTH})" - ) - if instructions and len(instructions) > MAX_INSTRUCTIONS_LENGTH: - raise ValueError( - f"Instructions exceed maximum length " - f"({len(instructions)} > {MAX_INSTRUCTIONS_LENGTH})" - ) - - -def _resolve_provider(provider: Optional[str], model: Optional[str]) -> dict[str, str]: - """Resolve provider and model, falling back to config defaults. - - Returns: - Dict with 'name' and 'model' keys. - - Raises: - ValueError: If neither explicit nor default values available. - """ - provider_name = provider or "" - model_name = model or "" - if not provider_name or not model_name: - inference = configuration.inference - provider_name = provider_name or inference.default_provider or "" - model_name = model_name or inference.default_model or "" - - if not provider_name or not model_name: - raise ValueError( - "Provider and model must be specified or configured as defaults " - "(inference.default_provider / inference.default_model)" - ) - return {"name": provider_name, "model": model_name} - - -def _get_or_create_runner( - *, - provider: dict[str, str], - instructions: Optional[str] = None, - mcp_server_names: Optional[list[str]] = None, -) -> ChatWorkflowRunner: - """Get or create the ChatWorkflowRunner singleton. - - Parameters: - provider: Resolved provider dict. - instructions: System prompt. - mcp_server_names: MCP server names to resolve. - - Returns: - ChatWorkflowRunner instance. - """ - global _runner - - if _runner is not None: - return _runner - - mcp_servers = resolve_mcp_servers(mcp_server_names) - - from utils.prompts import ( - get_system_prompt, - ) # pylint: disable=import-outside-toplevel - - resolved_instructions = get_system_prompt(instructions) - - config = ChatWorkflowConfig( - provider=provider, - system_prompt=resolved_instructions, - mcp_servers=mcp_servers or None, - ) - - try: - run_store = WorkflowStorageFactory.get_run_state_store() - transcript_store = WorkflowStorageFactory.get_transcript_store() - except RuntimeError: - run_store = None - transcript_store = None - - from workflow.middleware import ( - get_default_middleware, - ) # pylint: disable=import-outside-toplevel - - _runner = ChatWorkflowRunner( - run_store=run_store, - transcript_store=transcript_store, - config=config, - middlewares=get_default_middleware(), - ) - return _runner - - -def reset_runner() -> None: - """Reset the runner singleton (for testing).""" - global _runner - _runner = None - - -async def execute_query_via_direct_executor( # pylint: disable=too-many-arguments - *, - prompt: str, - model: Optional[str] = None, - provider: Optional[str] = None, - instructions: Optional[str] = None, - mcp_server_names: Optional[list[str]] = None, - output_schema: Optional[dict[str, Any]] = None, - context: Optional[dict[str, Any]] = None, - conversation_id: Optional[str] = None, - user_id: str = "", - username: str = "", -) -> StepResult: - """Execute a query using ChatWorkflowRunner (blocking). - - Parameters: - prompt: User's query text. - model: Model name (e.g. "gpt-4o-mini"). - provider: Provider name (e.g. "openai"). - instructions: System prompt / instructions. - mcp_server_names: MCP server names to include (None = all). - output_schema: Optional structured output schema. - context: Prior conversation context. - conversation_id: Conversation ID for multi-turn. - user_id: User identifier for audit logging. - username: Username for audit logging. - - Returns: - StepResult with agent response, transcript, and metrics. - - Raises: - ValueError: On validation failure. - """ - _validate_prompt(prompt, instructions) - if output_schema: - raise ValueError("output_schema is not yet supported via ChatWorkflowRunner") - resolved_provider = _resolve_provider(provider, model) - user_label = username or user_id or "anonymous" - - runner = _get_or_create_runner( - provider=resolved_provider, - instructions=instructions, - mcp_server_names=mcp_server_names, - ) - - if not conversation_id: - conversation_id = await runner.start({"user_id": user_id or None}) - - logger.info( - "Query via ChatWorkflowRunner: user=%s conv=%s model=%s:%s", - user_label, - conversation_id, - resolved_provider["name"], - resolved_provider["model"], - ) - - result = await runner.send_message(conversation_id, prompt) - - logger.info( - "Query completed: user=%s conv=%s status=%s duration_ms=%d " - "tokens_in=%d tokens_out=%d", - user_label, - conversation_id, - result.status, - result.duration_ms, - result.input_tokens, - result.output_tokens, - ) - - return result - - -async def stream_query_via_direct_executor( - *, - prompt: str, - model: Optional[str] = None, - provider: Optional[str] = None, - instructions: Optional[str] = None, - mcp_server_names: Optional[list[str]] = None, - output_schema: Optional[dict[str, Any]] = None, - context: Optional[dict[str, Any]] = None, - conversation_id: Optional[str] = None, - user_id: str = "", - username: str = "", -) -> AsyncIterator[StreamEvent]: - """Stream a query using ChatWorkflowRunner. - - Parameters: - prompt: User's query text. - model: Model name (e.g. "gpt-4o-mini"). - provider: Provider name (e.g. "openai"). - instructions: System prompt / instructions. - mcp_server_names: MCP server names to include (None = all). - output_schema: Not yet supported — raises ValueError if set. - context: Prior conversation context. - conversation_id: Conversation ID for multi-turn. - user_id: User identifier for audit logging. - username: Username for audit logging. - - Yields: - StreamEvent instances. - - Raises: - ValueError: On validation failure (raised BEFORE streaming starts). - """ - _validate_prompt(prompt, instructions) - if output_schema: - raise ValueError( - "output_schema is not yet supported via ChatWorkflowRunner streaming" - ) - resolved_provider = _resolve_provider(provider, model) - user_label = username or user_id or "anonymous" - - runner = _get_or_create_runner( - provider=resolved_provider, - instructions=instructions, - mcp_server_names=mcp_server_names, - ) - - if not conversation_id: - conversation_id = await runner.start({"user_id": user_id or None}) - - logger.info( - "Streaming query via ChatWorkflowRunner: user=%s conv=%s model=%s:%s", - user_label, - conversation_id, - resolved_provider["name"], - resolved_provider["model"], - ) - - async for event in runner.send_message_stream(conversation_id, prompt): - yield event - - logger.info("Stream completed: user=%s conv=%s", user_label, conversation_id) diff --git a/tests/e2e/cloud_agents/jaeger_helpers.py b/tests/e2e/cloud_agents/jaeger_helpers.py index 8e9fbf895..7549af963 100644 --- a/tests/e2e/cloud_agents/jaeger_helpers.py +++ b/tests/e2e/cloud_agents/jaeger_helpers.py @@ -1,8 +1,7 @@ """Shared Jaeger query helpers for cloud_agents E2E OTEL tracing tests. -Split out because both test_otel_tracing_e2e.py (ChatWorkflowRunner) and -test_workflow_tracing_e2e.py (LocalWorkflowRunner) need identical Jaeger -availability/query helpers. The `tracer` fixture itself (which installs a +Split out of test_workflow_tracing_e2e.py (LocalWorkflowRunner) so Jaeger +availability/query helpers can be reused. The `tracer` fixture itself (which installs a process-global TracerProvider) is intentionally NOT shared here — each file keeps its own copy and assumes standalone execution (see each file's docstring), since OpenTelemetry only allows one TracerProvider install per diff --git a/tests/e2e/cloud_agents/test_otel_tracing_e2e.py b/tests/e2e/cloud_agents/test_otel_tracing_e2e.py deleted file mode 100644 index dd317a1cc..000000000 --- a/tests/e2e/cloud_agents/test_otel_tracing_e2e.py +++ /dev/null @@ -1,223 +0,0 @@ -"""E2E tests for OTEL tracing through the agent execution path. - -Requires: -- OPENAI_API_KEY environment variable -- OTEL_ANONYMIZATION_SECRET environment variable (any value; see - docker-compose.yaml for the dev default) — the tracing middleware - refuses to run without it since it HMACs user_id before adding it - as a span attribute -- Jaeger running at localhost:4317 (OTLP) and localhost:16686 (query API) -- PostgreSQL reachable at WORKFLOW_PG_* (default: localhost:5432/lightspeed, - see docker-compose-harness.yaml); ChatWorkflowRunner requires a real - RunStateStore/TranscriptStore, which the fixture below sets up - (schema is migrated automatically on connect via Alembic) - -Usage: - OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4317 \ - OTEL_ANONYMIZATION_SECRET=lightspeed-stack-otel-anonymization-dev-default \ - uv run pytest tests/e2e/cloud_agents/test_otel_tracing_e2e.py -v -s -""" - -# pylint: disable=import-outside-toplevel,too-few-public-methods - -from __future__ import annotations - -import asyncio -import os -from collections.abc import AsyncIterator -from typing import Any - -import pytest - -from tests.e2e.cloud_agents.jaeger_helpers import ( - OTLP_ENDPOINT, - SERVICE_NAME, - check_jaeger_available, - query_jaeger_traces, -) - -pytestmark = [ - pytest.mark.skipif( - not os.environ.get("OPENAI_API_KEY"), - reason="OPENAI_API_KEY not set", - ), -] - -_PG_HOST = os.environ.get("WORKFLOW_PG_HOST", "localhost") -_PG_PORT = int(os.environ.get("WORKFLOW_PG_PORT", "5432")) -_PG_DB = os.environ.get("WORKFLOW_PG_DB", "lightspeed") -_PG_USER = os.environ.get("WORKFLOW_PG_USER", "lightspeed") -_PG_PASSWORD = os.environ.get("WORKFLOW_PG_PASSWORD", "lightspeed") - - -@pytest.fixture(name="workflow_storage", autouse=True) -async def workflow_storage_fixture() -> AsyncIterator[None]: - """Initialize WorkflowStorageFactory so ChatWorkflowRunner has a real store. - - execute_query_via_direct_executor() always routes through - ChatWorkflowRunner, which requires a connected RunStateStore / - TranscriptStore. Outside the FastAPI lifespan (main.py) nothing - initializes WorkflowStorageFactory, so it must be done here. - """ - from models.config import ( - PostgreSQLDatabaseConfiguration, - WorkflowEngineConfiguration, - ) - from workflow.query_executor import reset_runner - from workflow.storage import WorkflowStorageFactory - - pg_config = PostgreSQLDatabaseConfiguration( - host=_PG_HOST, - port=_PG_PORT, - db=_PG_DB, - user=_PG_USER, - password=_PG_PASSWORD, - ) - wf_config = WorkflowEngineConfiguration(enabled=True) - - reset_runner() - await WorkflowStorageFactory.initialize(pg_config, wf_config) - - yield - - reset_runner() - await WorkflowStorageFactory.cleanup() - - -@pytest.fixture(name="tracer", scope="module") -def tracer_fixture() -> Any: - """Set up OTEL tracer that exports to Jaeger.""" - from opentelemetry import trace - from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import ( - OTLPSpanExporter, - ) - from opentelemetry.sdk.resources import Resource - from opentelemetry.sdk.trace import TracerProvider - from opentelemetry.sdk.trace.export import SimpleSpanProcessor - - resource = Resource.create({"service.name": SERVICE_NAME}) - provider = TracerProvider(resource=resource) - exporter = OTLPSpanExporter(endpoint=OTLP_ENDPOINT, insecure=True) - provider.add_span_processor(SimpleSpanProcessor(exporter)) - trace.set_tracer_provider(provider) - - yield trace.get_tracer("e2e-test") - - provider.force_flush() - provider.shutdown() - - -class TestOtelTracingE2E: - """E2E tests verifying OTEL traces reach Jaeger.""" - - @pytest.mark.asyncio - async def test_query_direct_creates_trace(self, tracer: Any) -> None: - """A /query/direct call creates a trace visible in Jaeger.""" - if not await check_jaeger_available(): - pytest.skip("Jaeger not available") - - from configuration import configuration - from workflow.query_executor import execute_query_via_direct_executor - - configuration.init_from_dict( - { - "name": "otel-e2e-test", - "service": { - "host": "localhost", - "port": 8080, - "auth_enabled": False, - "workers": 1, - }, - "llama_stack": { - "use_as_library_client": False, - "url": "http://localhost:8321", - }, - "user_data_collection": {"feedback_enabled": False}, - "authentication": {"module": "noop"}, - } - ) - - with tracer.start_as_current_span("e2e.query_direct") as span: - span.set_attribute("test.name", "otel_tracing_e2e") - result = await execute_query_via_direct_executor( - prompt="What is 1+1? Reply with just the number.", - provider="openai", - model="gpt-4o-mini", - user_id="e2e-test-user", - ) - span.set_attribute("result.status", result.status) - span.set_attribute("result.tokens_in", result.input_tokens) - span.set_attribute("result.tokens_out", result.output_tokens) - trace_id = format(span.get_span_context().trace_id, "032x") - - assert result.status == "completed" - - await asyncio.sleep(2) - - traces = await query_jaeger_traces( - service=SERVICE_NAME, - operation="e2e.query_direct", - ) - - assert len(traces) >= 1, ( - f"Expected trace in Jaeger for service={SERVICE_NAME}, " - f"operation=e2e.query_direct. Found {len(traces)} traces." - ) - - found_trace = None - for t in traces: - if t.get("traceID") == trace_id: - found_trace = t - break - - assert found_trace is not None, ( - f"Trace {trace_id} not found in Jaeger. " - f"Available trace IDs: {[t.get('traceID') for t in traces]}" - ) - - spans = found_trace.get("spans", []) - assert len(spans) >= 1 - - root_span = next( - (s for s in spans if s.get("operationName") == "e2e.query_direct"), - None, - ) - assert root_span is not None - - tags = {t["key"]: t["value"] for t in root_span.get("tags", [])} - assert tags.get("test.name") == "otel_tracing_e2e" - assert tags.get("result.status") == "completed" - - @pytest.mark.asyncio - async def test_trace_includes_token_metrics(self, tracer: Any) -> None: - """Trace span includes token usage as attributes.""" - if not await check_jaeger_available(): - pytest.skip("Jaeger not available") - - from workflow.query_executor import execute_query_via_direct_executor - - with tracer.start_as_current_span("e2e.token_metrics") as span: - result = await execute_query_via_direct_executor( - prompt="Say hello.", - provider="openai", - model="gpt-4o-mini", - ) - span.set_attribute("llm.usage.input_tokens", result.input_tokens) - span.set_attribute("llm.usage.output_tokens", result.output_tokens) - span.set_attribute("llm.duration_ms", result.duration_ms) - trace_id = format(span.get_span_context().trace_id, "032x") - - await asyncio.sleep(2) - - traces = await query_jaeger_traces( - service=SERVICE_NAME, - operation="e2e.token_metrics", - ) - - found = next((t for t in traces if t.get("traceID") == trace_id), None) - assert found is not None - - root_span = found["spans"][0] - tags = {t["key"]: t["value"] for t in root_span.get("tags", [])} - assert tags.get("llm.usage.input_tokens", 0) > 0 - assert tags.get("llm.usage.output_tokens", 0) > 0 diff --git a/tests/e2e/cloud_agents/test_query_direct_handler_e2e.py b/tests/e2e/cloud_agents/test_query_direct_handler_e2e.py deleted file mode 100644 index e9d94c120..000000000 --- a/tests/e2e/cloud_agents/test_query_direct_handler_e2e.py +++ /dev/null @@ -1,95 +0,0 @@ -"""Handler-direct e2e tests for POST /v1/query/direct error paths. - -Calls `query_direct_handler.__wrapped__(...)` directly, bypassing the -`@authorize` decorator and FastAPI routing/request-validation entirely. -Only the error paths are covered here (the happy path is exercised via -the spawn=none step-executor tests in test_step_executor_e2e.py, since -/v1/query/direct's happy path is a thin wrapper around the same -DirectExecutor those tests already cover). - -No OPENAI_API_KEY needed -- these are pure validation-error paths that -raise before any provider/LLM call is made. - -Usage: - uv run pytest tests/e2e/cloud_agents/test_query_direct_handler_e2e.py -v -s -""" - -# pylint: disable=too-few-public-methods,unused-argument - -from __future__ import annotations - -from typing import Any - -import pytest -from fastapi import HTTPException - -from app.endpoints.query_direct import QueryDirectRequest, query_direct_handler -from configuration import configuration - -from .conftest import AUTH, make_request - -_CONFIG = { - "name": "e2e-query-direct-handler-test", - "service": { - "host": "localhost", - "port": 8080, - "auth_enabled": False, - "workers": 1, - }, - "llama_stack": { - "use_as_library_client": False, - "url": "http://localhost:8321", - }, - "user_data_collection": {"feedback_enabled": False}, - "authentication": {"module": "noop"}, -} - - -@pytest.fixture(name="e2e_config", scope="module") -def e2e_config_fixture() -> Any: - """Load config for handler-direct tests. No Llama Stack needed.""" - configuration.init_from_dict(_CONFIG) - return configuration - - -class TestQueryDirectErrorHandlingE2E: - """E2E tests for /v1/query/direct error paths.""" - - @pytest.mark.asyncio - async def test_missing_provider_returns_400(self, e2e_config: Any) -> None: - """Missing provider and model with no defaults returns 400.""" - body = QueryDirectRequest(query="Hello") - - with pytest.raises(HTTPException) as exc_info: - await query_direct_handler.__wrapped__(make_request(), body, AUTH) - assert exc_info.value.status_code == 400 - assert "Provider and model" in str(exc_info.value.detail) - - @pytest.mark.asyncio - async def test_prompt_too_long_returns_400(self, e2e_config: Any) -> None: - """Oversized prompt returns 400.""" - body = QueryDirectRequest( - query="x" * 200_000, - provider="openai", - model="gpt-4o-mini", - ) - - with pytest.raises(HTTPException) as exc_info: - await query_direct_handler.__wrapped__(make_request(), body, AUTH) - assert exc_info.value.status_code == 400 - assert "maximum length" in str(exc_info.value.detail) - - @pytest.mark.asyncio - async def test_unknown_mcp_server_returns_400(self, e2e_config: Any) -> None: - """Unknown MCP server name returns 400.""" - body = QueryDirectRequest( - query="Hello", - provider="openai", - model="gpt-4o-mini", - mcp_servers=["nonexistent-server"], - ) - - with pytest.raises(HTTPException) as exc_info: - await query_direct_handler.__wrapped__(make_request(), body, AUTH) - assert exc_info.value.status_code == 400 - assert "Unknown MCP server" in str(exc_info.value.detail) diff --git a/tests/e2e/cloud_agents/test_step_executor_e2e.py b/tests/e2e/cloud_agents/test_step_executor_e2e.py index a188e6c67..fd595e7f4 100644 --- a/tests/e2e/cloud_agents/test_step_executor_e2e.py +++ b/tests/e2e/cloud_agents/test_step_executor_e2e.py @@ -342,20 +342,3 @@ async def test_result_has_transcript(self) -> None: assert result.status == "completed" assert isinstance(result.transcript, list) assert len(result.transcript) >= 1 - - -class TestPromptValidation: - """Validation helpers used by the step-executor path. - - Note: this doesn't exercise the step-executor itself, just the - private `_validate_prompt()` function it (and query_direct_handler) - call before dispatching -- arguably a tests/unit candidate rather - than e2e, kept here as-is pending a separate decision on moving it. - """ - - def test_stream_validation_error(self) -> None: - """Validation helpers raise on invalid input.""" - from workflow.query_executor import _validate_prompt - - with pytest.raises(ValueError, match="maximum length"): - _validate_prompt("x" * 200_000) diff --git a/tests/e2e/cloud_agents/test_workflow_tracing_e2e.py b/tests/e2e/cloud_agents/test_workflow_tracing_e2e.py index 72d8e0876..1d5c4193d 100644 --- a/tests/e2e/cloud_agents/test_workflow_tracing_e2e.py +++ b/tests/e2e/cloud_agents/test_workflow_tracing_e2e.py @@ -1,7 +1,6 @@ """E2E tests verifying OTEL trace chaining for LocalWorkflowRunner runs. -Companion to test_otel_tracing_e2e.py (which covers ChatWorkflowRunner) — -this file covers the workflow-runner path (LocalWorkflowRunner, i.e. +This file covers the workflow-runner path (LocalWorkflowRunner, i.e. /v1/workflows/*). See jameswnl/lightspeed-stack#20 for full context and jameswnl/lightspeed-cloud-agents#179 for the originating cloud-agents issue (split into #181 and #183, both now resolved). diff --git a/tests/unit/app/test_routers.py b/tests/unit/app/test_routers.py index d8c8917f1..36ab7984c 100644 --- a/tests/unit/app/test_routers.py +++ b/tests/unit/app/test_routers.py @@ -123,7 +123,7 @@ def test_include_routers() -> None: include_routers(app) # are all routers added? - assert len(app.routers) == 29 + assert len(app.routers) == 28 assert root.router in app.get_routers() assert info.router in app.get_routers() assert models.router in app.get_routers() @@ -166,7 +166,7 @@ def test_check_prefixes() -> None: include_routers(app) # are all routers added? - assert len(app.routers) == 29 + assert len(app.routers) == 28 assert app.get_router_prefix(root.router) == "" assert app.get_router_prefix(info.router) == "/v1" assert app.get_router_prefix(models.router) == "/v1" diff --git a/tests/unit/cloud_agents/test_middleware.py b/tests/unit/cloud_agents/test_middleware.py deleted file mode 100644 index 6272f9958..000000000 --- a/tests/unit/cloud_agents/test_middleware.py +++ /dev/null @@ -1,185 +0,0 @@ -"""Unit tests for lightspeed-stack StepMiddleware implementations.""" - -from __future__ import annotations - -import pytest -from pytest_mock import MockerFixture - -from utils.otel_tracing import SpanAttributes -from workflow.middleware import ( - AuditMiddleware, - QuotaMiddleware, - TracingMiddleware, - get_default_middleware, -) - - -def _make_step_input(mocker: MockerFixture, user_id: str = "test-user") -> object: - """Create a mock StepInput.""" - step_input = mocker.MagicMock() - step_input.step_name = "test-step" - step_input.provider = {"name": "openai", "model": "gpt-4o-mini"} - step_input.metadata = mocker.MagicMock() - step_input.metadata.user_id = user_id - return step_input - - -def _make_result(mocker: MockerFixture) -> object: - """Create a mock StepResult.""" - result = mocker.MagicMock() - result.status = "completed" - result.input_tokens = 50 - result.output_tokens = 25 - result.duration_ms = 500 - return result - - -class TestTracingMiddleware: - """Tests for TracingMiddleware.""" - - @pytest.mark.asyncio - async def test_before_passes_through(self, mocker: MockerFixture) -> None: - """Before hook returns step_input unchanged.""" - mw = TracingMiddleware() - step_input = _make_step_input(mocker) - result = await mw.before(step_input) - assert result is step_input - - @pytest.mark.asyncio - async def test_after_sets_span_attributes(self, mocker: MockerFixture) -> None: - """After hook sets OTEL attributes using semantic conventions.""" - mock_span = mocker.MagicMock() - mock_span.is_recording.return_value = True - mocker.patch( - "workflow.middleware.trace.get_current_span", - return_value=mock_span, - ) - mocker.patch( - "workflow.middleware.anonymize_value", - side_effect=lambda x: f"anon:{x}", - ) - - mw = TracingMiddleware() - step_input = _make_step_input(mocker) - result = _make_result(mocker) - - returned = await mw.after(step_input, result) - - assert returned is result - mock_span.set_attribute.assert_any_call( - SpanAttributes.LLM_MODEL_ID, "openai:gpt-4o-mini" - ) - mock_span.set_attribute.assert_any_call( - SpanAttributes.LLM_USAGE_INPUT_TOKENS, 50 - ) - mock_span.set_attribute.assert_any_call( - SpanAttributes.LLM_USAGE_OUTPUT_TOKENS, 25 - ) - mock_span.set_attribute.assert_any_call( - SpanAttributes.USER_ID, "anon:test-user" - ) - - @pytest.mark.asyncio - async def test_after_noop_when_not_recording(self, mocker: MockerFixture) -> None: - """No-op when span is not recording.""" - mock_span = mocker.MagicMock() - mock_span.is_recording.return_value = False - mocker.patch( - "workflow.middleware.trace.get_current_span", - return_value=mock_span, - ) - - mw = TracingMiddleware() - await mw.after(_make_step_input(mocker), _make_result(mocker)) - - mock_span.set_attribute.assert_not_called() - - @pytest.mark.asyncio - async def test_after_handles_none_metadata(self, mocker: MockerFixture) -> None: - """Handles step_input with metadata=None.""" - mock_span = mocker.MagicMock() - mock_span.is_recording.return_value = True - mocker.patch( - "workflow.middleware.trace.get_current_span", - return_value=mock_span, - ) - - mw = TracingMiddleware() - step_input = _make_step_input(mocker) - step_input.metadata = None - result = _make_result(mocker) - - await mw.after(step_input, result) - - calls = [str(c) for c in mock_span.set_attribute.call_args_list] - assert not any("user.id" in c for c in calls) - - @pytest.mark.asyncio - async def test_after_skips_empty_provider(self, mocker: MockerFixture) -> None: - """Skips model_id attribute when provider is empty.""" - mock_span = mocker.MagicMock() - mock_span.is_recording.return_value = True - mocker.patch( - "workflow.middleware.trace.get_current_span", - return_value=mock_span, - ) - - mw = TracingMiddleware() - step_input = _make_step_input(mocker) - step_input.provider = {} - result = _make_result(mocker) - - await mw.after(step_input, result) - - calls = [str(c) for c in mock_span.set_attribute.call_args_list] - assert not any("llm.model.id" in c for c in calls) - - -class TestAuditMiddleware: - """Tests for AuditMiddleware.""" - - @pytest.mark.asyncio - async def test_passes_through(self, mocker: MockerFixture) -> None: - """Audit middleware passes input and result through.""" - mw = AuditMiddleware() - step_input = _make_step_input(mocker) - result = _make_result(mocker) - - assert await mw.before(step_input) is step_input - assert await mw.after(step_input, result) is result - - @pytest.mark.asyncio - async def test_handles_none_metadata(self, mocker: MockerFixture) -> None: - """Uses 'anonymous' when metadata is None.""" - mw = AuditMiddleware() - step_input = _make_step_input(mocker) - step_input.metadata = None - - returned = await mw.before(step_input) - assert returned is step_input - - -class TestQuotaMiddleware: - """Tests for QuotaMiddleware.""" - - @pytest.mark.asyncio - async def test_passes_through(self, mocker: MockerFixture) -> None: - """Quota middleware passes input and result through.""" - mw = QuotaMiddleware() - step_input = _make_step_input(mocker) - result = _make_result(mocker) - - assert await mw.before(step_input) is step_input - assert await mw.after(step_input, result) is result - - -class TestGetDefaultMiddleware: - """Tests for get_default_middleware.""" - - def test_returns_three_middlewares(self) -> None: - """Default stack has tracing, audit, quota.""" - middlewares = get_default_middleware() - assert len(middlewares) == 3 - assert isinstance(middlewares[0], TracingMiddleware) - assert isinstance(middlewares[1], AuditMiddleware) - assert isinstance(middlewares[2], QuotaMiddleware) diff --git a/tests/unit/cloud_agents/test_query_executor.py b/tests/unit/cloud_agents/test_query_executor.py deleted file mode 100644 index af8715048..000000000 --- a/tests/unit/cloud_agents/test_query_executor.py +++ /dev/null @@ -1,289 +0,0 @@ -"""Unit tests for the query executor bridge.""" - -# pylint: disable=protected-access - -from __future__ import annotations - -import pytest -from pytest_mock import MockerFixture - -from workflow.query_executor import ( - _resolve_provider, - _validate_prompt, - execute_query_via_direct_executor, - resolve_mcp_servers, - stream_query_via_direct_executor, -) - - -class TestResolveMcpServers: - """Tests for resolve_mcp_servers.""" - - def test_resolves_all_servers(self, mocker: MockerFixture) -> None: - """Resolves all configured MCP servers when no filter.""" - mock_server1 = mocker.MagicMock() - mock_server1.name = "kubectl" - mock_server1.url = "http://mcp-kubectl:8080/sse" - mock_server1.resolved_authorization_headers = {} - - mock_server2 = mocker.MagicMock() - mock_server2.name = "github" - mock_server2.url = "http://mcp-github:8080/sse" - mock_server2.resolved_authorization_headers = { - "Authorization": "Bearer token123" - } - - mock_config = mocker.patch("workflow.query_executor.configuration") - mock_config.mcp_servers = [mock_server1, mock_server2] - - result = resolve_mcp_servers() - - assert len(result) == 2 - assert result[0] == {"name": "kubectl", "url": "http://mcp-kubectl:8080/sse"} - assert result[1]["headers"]["Authorization"] == "Bearer token123" - - def test_filters_by_name(self, mocker: MockerFixture) -> None: - """Only resolves named servers when filter is provided.""" - mock_server1 = mocker.MagicMock() - mock_server1.name = "kubectl" - mock_server1.url = "http://mcp-kubectl:8080/sse" - mock_server1.resolved_authorization_headers = {} - - mock_server2 = mocker.MagicMock() - mock_server2.name = "github" - mock_server2.url = "http://mcp-github:8080/sse" - mock_server2.resolved_authorization_headers = {} - - mock_config = mocker.patch("workflow.query_executor.configuration") - mock_config.mcp_servers = [mock_server1, mock_server2] - - result = resolve_mcp_servers(server_names=["kubectl"]) - - assert len(result) == 1 - assert result[0]["name"] == "kubectl" - - def test_empty_when_no_servers(self, mocker: MockerFixture) -> None: - """Returns empty list when no servers configured.""" - mock_config = mocker.patch("workflow.query_executor.configuration") - mock_config.mcp_servers = [] - - result = resolve_mcp_servers() - assert result == [] - - def test_unknown_server_raises(self, mocker: MockerFixture) -> None: - """Unknown server name raises ValueError.""" - mock_server = mocker.MagicMock() - mock_server.name = "kubectl" - mock_config = mocker.patch("workflow.query_executor.configuration") - mock_config.mcp_servers = [mock_server] - - with pytest.raises(ValueError, match="Unknown MCP server"): - resolve_mcp_servers(server_names=["nonexistent"]) - - -class TestValidatePrompt: - """Tests for _validate_prompt.""" - - def test_valid_prompt(self) -> None: - """Normal prompt passes validation.""" - _validate_prompt("Hello, world!") - - def test_prompt_too_long_raises(self) -> None: - """Oversized prompt raises ValueError.""" - with pytest.raises(ValueError, match="exceeds maximum length"): - _validate_prompt("x" * 200_000) - - def test_instructions_too_long_raises(self) -> None: - """Oversized instructions raises ValueError.""" - with pytest.raises(ValueError, match="Instructions exceed"): - _validate_prompt("Hello", "x" * 100_000) - - -class TestResolveProvider: - """Tests for _resolve_provider.""" - - def test_explicit_values(self, mocker: MockerFixture) -> None: - """Returns explicit provider and model.""" - mocker.patch("workflow.query_executor.configuration") - result = _resolve_provider("openai", "gpt-4o-mini") - assert result == {"name": "openai", "model": "gpt-4o-mini"} - - def test_falls_back_to_defaults(self, mocker: MockerFixture) -> None: - """Falls back to config defaults.""" - mock_config = mocker.patch("workflow.query_executor.configuration") - mock_config.inference.default_provider = "openai" - mock_config.inference.default_model = "gpt-4o-mini" - - result = _resolve_provider(None, None) - assert result == {"name": "openai", "model": "gpt-4o-mini"} - - def test_missing_both_raises(self, mocker: MockerFixture) -> None: - """Raises ValueError when no provider/model and no defaults.""" - mock_config = mocker.patch("workflow.query_executor.configuration") - mock_config.inference.default_provider = None - mock_config.inference.default_model = None - - with pytest.raises(ValueError, match="Provider and model must be"): - _resolve_provider(None, None) - - -class TestExecuteQueryViaChatWorkflowRunner: - """Tests for execute_query_via_direct_executor.""" - - @pytest.mark.asyncio - async def test_executes_query(self, mocker: MockerFixture) -> None: - """Sends message to ChatWorkflowRunner and returns result.""" - mock_result = mocker.MagicMock() - mock_result.status = "completed" - mock_result.output = {"response": "Hello"} - mock_result.input_tokens = 10 - mock_result.output_tokens = 5 - mock_result.duration_ms = 500 - - mock_runner = mocker.AsyncMock() - mock_runner.start.return_value = "chat-abc123" - mock_runner.send_message.return_value = mock_result - - mocker.patch( - "workflow.query_executor._get_or_create_runner", - return_value=mock_runner, - ) - mocker.patch("workflow.query_executor.configuration").mcp_servers = [] - mocker.patch( - "workflow.query_executor._resolve_provider", - return_value={"name": "openai", "model": "gpt-4o-mini"}, - ) - - result = await execute_query_via_direct_executor( - prompt="Hello", - provider="openai", - model="gpt-4o-mini", - ) - - assert result.status == "completed" - mock_runner.start.assert_called_once() - mock_runner.send_message.assert_called_once_with("chat-abc123", "Hello") - - @pytest.mark.asyncio - async def test_reuses_conversation_id(self, mocker: MockerFixture) -> None: - """Uses provided conversation_id instead of starting new.""" - mock_result = mocker.MagicMock() - mock_result.status = "completed" - mock_result.input_tokens = 10 - mock_result.output_tokens = 5 - mock_result.duration_ms = 500 - - mock_runner = mocker.AsyncMock() - mock_runner.send_message.return_value = mock_result - - mocker.patch( - "workflow.query_executor._get_or_create_runner", - return_value=mock_runner, - ) - mocker.patch("workflow.query_executor.configuration").mcp_servers = [] - mocker.patch( - "workflow.query_executor._resolve_provider", - return_value={"name": "openai", "model": "gpt-4o-mini"}, - ) - - await execute_query_via_direct_executor( - prompt="Hello", - provider="openai", - model="gpt-4o-mini", - conversation_id="existing-conv", - ) - - mock_runner.start.assert_not_called() - mock_runner.send_message.assert_called_once_with("existing-conv", "Hello") - - @pytest.mark.asyncio - async def test_output_schema_raises(self, mocker: MockerFixture) -> None: - """output_schema raises ValueError (not yet supported).""" - mocker.patch("workflow.query_executor.configuration").mcp_servers = [] - mocker.patch( - "workflow.query_executor._resolve_provider", - return_value={"name": "openai", "model": "gpt-4o-mini"}, - ) - - with pytest.raises(ValueError, match="not yet supported"): - await execute_query_via_direct_executor( - prompt="Hello", - provider="openai", - model="gpt-4o-mini", - output_schema={"type": "object"}, - ) - - -class TestStreamQueryViaChatWorkflowRunner: - """Tests for stream_query_via_direct_executor.""" - - @pytest.mark.asyncio - async def test_yields_stream_events(self, mocker: MockerFixture) -> None: - """Yields events from ChatWorkflowRunner.send_message_stream.""" - from cloud_agents.workflow.executor.step.base import StreamEvent - - mock_events = [ - StreamEvent(type="token", data={"delta": "Hello"}), - StreamEvent(type="complete", data={}, result=mocker.MagicMock()), - ] - - async def mock_stream(_wf_id: str, _prompt: str): # type: ignore[no-untyped-def] - for event in mock_events: - yield event - - mock_runner = mocker.AsyncMock() - mock_runner.start.return_value = "chat-stream-1" - mock_runner.send_message_stream = mock_stream - - mocker.patch( - "workflow.query_executor._get_or_create_runner", - return_value=mock_runner, - ) - mocker.patch("workflow.query_executor.configuration").mcp_servers = [] - mocker.patch( - "workflow.query_executor._resolve_provider", - return_value={"name": "openai", "model": "gpt-4o-mini"}, - ) - - events = [] - async for event in stream_query_via_direct_executor( - prompt="Hello", - provider="openai", - model="gpt-4o-mini", - ): - events.append(event) - - assert len(events) == 2 - assert events[0].type == "token" - assert events[1].type == "complete" - - @pytest.mark.asyncio - async def test_reuses_conversation_id(self, mocker: MockerFixture) -> None: - """Streaming with conversation_id skips start().""" - from cloud_agents.workflow.executor.step.base import StreamEvent - - async def mock_stream(_wf_id: str, _prompt: str): # type: ignore[no-untyped-def] - yield StreamEvent(type="complete", data={}, result=mocker.MagicMock()) - - mock_runner = mocker.AsyncMock() - mock_runner.send_message_stream = mock_stream - - mocker.patch( - "workflow.query_executor._get_or_create_runner", - return_value=mock_runner, - ) - mocker.patch("workflow.query_executor.configuration").mcp_servers = [] - mocker.patch( - "workflow.query_executor._resolve_provider", - return_value={"name": "openai", "model": "gpt-4o-mini"}, - ) - - async for _ in stream_query_via_direct_executor( - prompt="Hello", - provider="openai", - model="gpt-4o-mini", - conversation_id="existing-stream-conv", - ): - pass - - mock_runner.start.assert_not_called()