LCORE-3312: Refactored exception handling

Makefile

@@ -144,7 +144,7 @@ clean-llama-stack: remove-llama-stack-container ## Remove container and image
 
 run-llama-stack: ## Start Llama Stack with enriched config (for local service mode)
 	uv run src/llama_stack_configuration.py -c $(CONFIG) -i $(LLAMA_STACK_CONFIG) -o $(LLAMA_STACK_CONFIG) && \
-	uv run ogx stack run $(LLAMA_STACK_CONFIG)
+	uv run ogx stack run --insecure $(LLAMA_STACK_CONFIG)
 
 test-unit: ## Run the unit tests
 	@echo "Running unit tests..."

pyproject.toml

@@ -80,7 +80,7 @@ dependencies = [
     # Used for token estimation before LLM calls (LCORE-1569 / conversation compaction)
     "tiktoken>=0.8.0",
     # Used for Pydantic AI
-    "pydantic-ai==2.16.0",
+    "pydantic-ai>=2.23.0",
     "pydantic-ai-skills>=0.11.0",
     # Used for OpenTelemetry instrumentation
     "opentelemetry-distro>=0.49b0",

run.yaml

@@ -59,6 +59,8 @@ providers:
     provider_type: inline::reference
 server:
   port: 8321
+  # OGX 1.3+ requires TLS by default; e2e/local use plain HTTP.
+  insecure: true
 storage:
   backends:
     kv_default: # Define the storage backend type for RAG, in this case registry and RAG are unified i.e. information on registered resources (e.g. models, vector_stores) are saved together with the RAG chunks

scripts/llama-stack-entrypoint.sh

@@ -19,9 +19,10 @@ if [ -f "$LIGHTSPEED_CONFIG" ]; then
 
     if [ -f "$ENRICHED_CONFIG" ] && [ "$ENRICHMENT_FAILED" -eq 0 ]; then
         echo "Using enriched config: $ENRICHED_CONFIG"
-        exec ogx stack run "$ENRICHED_CONFIG"
+        # OGX 1.3+ requires TLS unless --insecure is set (e2e/local HTTP).
+        exec ogx stack run --insecure "$ENRICHED_CONFIG"
     fi
 fi
 
 echo "Using original config: $INPUT_CONFIG"
-exec ogx stack run "$INPUT_CONFIG"
+exec ogx stack run --insecure "$INPUT_CONFIG"

src/app/endpoints/a2a.py

@@ -33,7 +33,7 @@
 )
 from a2a.utils import new_agent_text_message, new_task
 from fastapi import APIRouter, Depends, HTTPException, Request, status
-from ogx_client import APIConnectionError, APIStatusError
+from ogx_client import ApiException
 from pydantic_ai import AgentRunResultEvent
 from pydantic_ai.exceptions import AgentRunError
 from pydantic_ai.messages import (
@@ -378,7 +378,7 @@ async def _process_task_streaming(  # pylint: disable=too-many-locals
                 configuration,
                 shields=query_request.shield_ids,
             )
-        except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as e:
+        except (AgentRunError, ApiException, RuntimeError) as e:
             error_response = map_agent_inference_error(e, query_request.model or "")
             logger.error("Error preparing A2A agent: %s", str(e), exc_info=True)
             await task_updater.update_status(
@@ -437,7 +437,7 @@ async def _process_task_streaming(  # pylint: disable=too-many-locals
             ):
                 aggregator.process_event(a2a_event)
                 await event_queue.enqueue_event(a2a_event)
-        except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as e:
+        except (AgentRunError, ApiException, RuntimeError) as e:
             error_response = map_agent_inference_error(e, responses_params.model)
             logger.error("Error during A2A agent run: %s", str(e), exc_info=True)
             await task_updater.update_status(

src/app/endpoints/conversations_v1.py

@@ -4,10 +4,7 @@
 
 from fastapi import APIRouter, Depends, HTTPException, Request
 from ogx_api import ConversationNotFoundError, InvalidParameterError
-from ogx_client import (
-    APIConnectionError,
-    APIStatusError,
-)
+from ogx_client import ApiException
 from sqlalchemy.exc import SQLAlchemyError
 
 from app.database import get_session
@@ -273,15 +270,21 @@ async def get_conversation_endpoint_handler(  # pylint: disable=too-many-locals,
             chat_history=chat_history,
         )
 
-    except APIConnectionError as e:
-        logger.error("Unable to connect to Llama Stack: %s", e)
-        response = ServiceUnavailableResponse(
-            backend_name="OGX", cause=str(e)
+    except ApiException as e:
+        if not e.status:
+            logger.error("Unable to connect to Llama Stack: %s", e)
+            response = ServiceUnavailableResponse(
+                backend_name="OGX", cause=str(e)
+            ).model_dump()
+            raise HTTPException(**response) from e
+        # In library mode, ConversationNotFoundError is raised instead of ApiException
+        logger.error("Conversation not found: %s", e)
+        response = NotFoundResponse(
+            resource="conversation", resource_id=normalized_conv_id
         ).model_dump()
         raise HTTPException(**response) from e
-
-    except (APIStatusError, ConversationNotFoundError) as e:
-        # In library mode, ConversationNotFoundError is raised instead of APIStatusError
+    except ConversationNotFoundError as e:
+        # In library mode, ConversationNotFoundError is raised instead of ApiException
         logger.error("Conversation not found: %s", e)
         response = NotFoundResponse(
             resource="conversation", resource_id=normalized_conv_id
@@ -384,12 +387,17 @@ async def delete_conversation_endpoint_handler(
             delete_response.deleted,
         )
 
-    except APIConnectionError as e:
-        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
-        raise HTTPException(**response.model_dump()) from e
-
-    except (APIStatusError, ConversationNotFoundError, InvalidParameterError):
-        # In library mode, ConversationNotFoundError is raised instead of APIStatusError
+    except ApiException as e:
+        if not e.status:
+            response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
+            raise HTTPException(**response.model_dump()) from e
+        # In library mode, ConversationNotFoundError is raised instead of ApiException
+        logger.warning(
+            "Conversation %s in LlamaStack not found. Treating as already deleted.",
+            normalized_conv_id,
+        )
+    except (ConversationNotFoundError, InvalidParameterError):
+        # In library mode, ConversationNotFoundError is raised instead of ApiException
         logger.warning(
             "Conversation %s in LlamaStack not found. Treating as already deleted.",
             normalized_conv_id,
@@ -460,8 +468,7 @@ async def update_conversation_endpoint_handler(
 
     # If reached this, user is authorized to update this conversation
     try:
-        conversation = retrieve_conversation(normalized_conv_id)
-        if conversation is None:
+        if retrieve_conversation(normalized_conv_id) is None:
             response = NotFoundResponse(
                 resource="conversation", resource_id=normalized_conv_id
             ).model_dump()
@@ -520,14 +527,20 @@ async def update_conversation_endpoint_handler(
             message="Topic summary updated successfully",
         )
 
-    except APIConnectionError as e:
-        response = ServiceUnavailableResponse(
-            backend_name="OGX", cause=str(e)
+    except ApiException as e:
+        if not e.status:
+            response = ServiceUnavailableResponse(
+                backend_name="OGX", cause=str(e)
+            ).model_dump()
+            raise HTTPException(**response) from e
+        # In library mode, ConversationNotFoundError is raised instead of ApiException
+        logger.error("Conversation not found: %s", e)
+        response = NotFoundResponse(
+            resource="conversation", resource_id=normalized_conv_id
         ).model_dump()
         raise HTTPException(**response) from e
-
-    except (APIStatusError, ConversationNotFoundError) as e:
-        # In library mode, ConversationNotFoundError is raised instead of APIStatusError
+    except ConversationNotFoundError as e:
+        # In library mode, ConversationNotFoundError is raised instead of ApiException
         logger.error("Conversation not found: %s", e)
         response = NotFoundResponse(
             resource="conversation", resource_id=normalized_conv_id

src/app/endpoints/health.py

@@ -8,7 +8,7 @@
 from typing import Annotated, Any
 
 from fastapi import APIRouter, Depends, Response, status
-from ogx_client import APIConnectionError
+from ogx_client import ApiException
 
 from authentication import get_auth_dependency
 from authentication.interface import AuthTuple
@@ -78,7 +78,7 @@ async def get_providers_health_statuses() -> list[ProviderHealthStatus]:
             for provider in providers
         ]
 
-    except APIConnectionError as e:
+    except ApiException as e:
         logger.error("Failed to check providers health: %s", e)
         return [
             ProviderHealthStatus(

src/app/endpoints/info.py

@@ -3,7 +3,7 @@
 from typing import Annotated, Any
 
 from fastapi import APIRouter, Depends, HTTPException, Request
-from ogx_client import APIConnectionError
+from ogx_client import ApiException
 
 from authentication import get_auth_dependency
 from authentication.interface import AuthTuple
@@ -83,7 +83,7 @@ async def info_endpoint_handler(
             llama_stack_version=llama_stack_version,
         )
     # connection to Llama Stack server
-    except APIConnectionError as e:
+    except ApiException as e:
         logger.error("Unable to connect to Llama Stack: %s", e)
         response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
         raise HTTPException(**response.model_dump()) from e

src/app/endpoints/models.py

@@ -4,7 +4,7 @@
 
 from fastapi import APIRouter, HTTPException, Query, Request
 from fastapi.params import Depends
-from ogx_client import APIConnectionError
+from ogx_client import ApiException
 
 from authentication import get_auth_dependency
 from authentication.interface import AuthTuple
@@ -106,7 +106,7 @@ async def models_endpoint_handler(
         return ModelsResponse(models=parsed_models)
 
     # Connection to Llama Stack server failed
-    except APIConnectionError as e:
+    except ApiException as e:
         logger.error("Unable to connect to Llama Stack: %s", e)
         response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
         raise HTTPException(**response.model_dump()) from e

src/app/endpoints/prompts.py

@@ -3,9 +3,7 @@
 from typing import Annotated, Any, Optional
 
 from fastapi import APIRouter, Depends, HTTPException, Request
-from ogx_client import APIConnectionError, BadRequestError
-from ogx_client import APIStatusError as LLSApiStatusError
-from openai._exceptions import APIStatusError as OpenAIAPIStatusError
+from ogx_client import ApiException, BadRequestError
 
 from authentication import get_auth_dependency
 from authentication.interface import AuthTuple
@@ -139,11 +137,12 @@ async def create_prompt_handler(
         payload = body.model_dump(exclude_none=True)
         created = await client.prompts.create(**payload)
         return PromptResourceResponse.model_validate(created.model_dump())
-    except APIConnectionError as e:
-        logger.error("Unable to connect to Llama Stack: %s", e)
-        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
-        raise HTTPException(**response.model_dump()) from e
-    except (LLSApiStatusError, OpenAIAPIStatusError) as e:
+    except ApiException as e:
+        if not e.status:
+            logger.error("Unable to connect to Llama Stack: %s", e)
+            response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
+            raise HTTPException(**response.model_dump()) from e
+
         logger.error("API status error while creating prompt: %s", e)
         error_response = handle_known_apistatus_errors(e, "llama-stack")
         raise HTTPException(**error_response.model_dump()) from e
@@ -188,11 +187,12 @@ async def list_prompts_handler(
         items = await client.prompts.list()
         data = [PromptResourceResponse.model_validate(p.model_dump()) for p in items]
         return PromptsListResponse(data=data)
-    except APIConnectionError as e:
-        logger.error("Unable to connect to Llama Stack: %s", e)
-        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
-        raise HTTPException(**response.model_dump()) from e
-    except (LLSApiStatusError, OpenAIAPIStatusError) as e:
+    except ApiException as e:
+        if not e.status:
+            logger.error("Unable to connect to Llama Stack: %s", e)
+            response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
+            raise HTTPException(**response.model_dump()) from e
+
         logger.error("API status error while listing prompts: %s", e)
         error_response = handle_known_apistatus_errors(e, "llama-stack")
         raise HTTPException(**error_response.model_dump()) from e
@@ -247,15 +247,16 @@ async def get_prompt_handler(
         client = AsyncOgxClientHolder().get_client()
         retrieved = await client.prompts.retrieve(prompt_id, version=version)
         return PromptResourceResponse.model_validate(retrieved.model_dump())
-    except APIConnectionError as e:
-        logger.error("Unable to connect to Llama Stack: %s", e)
-        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
-        raise HTTPException(**response.model_dump()) from e
     except (BadRequestError, ValueError) as e:
         logger.error("Prompt not found: %s", e)
         response = NotFoundResponse(resource="prompt", resource_id=prompt_id)
         raise HTTPException(**response.model_dump()) from e
-    except (LLSApiStatusError, OpenAIAPIStatusError) as e:
+    except ApiException as e:
+        if not e.status:
+            logger.error("Unable to connect to Llama Stack: %s", e)
+            response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
+            raise HTTPException(**response.model_dump()) from e
+
         logger.error("API status error while retrieving prompt: %s", e)
         error_response = handle_known_apistatus_errors(e, "llama-stack")
         raise HTTPException(**error_response.model_dump()) from e
@@ -316,15 +317,16 @@ async def update_prompt_handler(
         payload = body.model_dump(exclude_none=True, exclude_unset=True)
         updated = await client.prompts.update(prompt_id, **payload)
         return PromptResourceResponse.model_validate(updated.model_dump())
-    except APIConnectionError as e:
-        logger.error("Unable to connect to Llama Stack: %s", e)
-        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
-        raise HTTPException(**response.model_dump()) from e
     except (BadRequestError, ValueError) as e:
         logger.error("Prompt update failed: %s", e)
         response = NotFoundResponse(resource="prompt", resource_id=prompt_id)
         raise HTTPException(**response.model_dump()) from e
-    except (LLSApiStatusError, OpenAIAPIStatusError) as e:
+    except ApiException as e:
+        if not e.status:
+            logger.error("Unable to connect to Llama Stack: %s", e)
+            response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
+            raise HTTPException(**response.model_dump()) from e
+
         logger.error("API status error while updating prompt: %s", e)
         error_response = handle_known_apistatus_errors(e, "llama-stack")
         raise HTTPException(**error_response.model_dump()) from e
@@ -381,14 +383,15 @@ async def delete_prompt_handler(
         client = AsyncOgxClientHolder().get_client()
         await client.prompts.delete(prompt_id)
         return PromptDeleteResponse(deleted=True, prompt_id=prompt_id)
-    except APIConnectionError as e:
-        logger.error("Unable to connect to Llama Stack: %s", e)
-        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
-        raise HTTPException(**response.model_dump()) from e
     except (BadRequestError, ValueError) as e:
         logger.error("Prompt delete failed: %s", e)
         return PromptDeleteResponse(deleted=False, prompt_id=prompt_id)
-    except (LLSApiStatusError, OpenAIAPIStatusError) as e:
+    except ApiException as e:
+        if not e.status:
+            logger.error("Unable to connect to Llama Stack: %s", e)
+            response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
+            raise HTTPException(**response.model_dump()) from e
+
         logger.error("API status error while deleting prompt: %s", e)
         error_response = handle_known_apistatus_errors(e, "llama-stack")
         raise HTTPException(**error_response.model_dump()) from e

src/app/endpoints/providers.py

@@ -4,7 +4,7 @@
 
 from fastapi import APIRouter, HTTPException, Request
 from fastapi.params import Depends
-from ogx_client import APIConnectionError, BadRequestError
+from ogx_client import ApiException, BadRequestError
 from ogx_client.models.list_providers_response import ListProvidersResponse
 
 from authentication import get_auth_dependency
@@ -92,15 +92,17 @@ async def providers_endpoint_handler(
     try:
         client = AsyncOgxClientHolder().get_client()
         providers: ListProvidersResponse = await client.providers.list()
-    except APIConnectionError as e:
+    except ApiException as e:
         logger.error("Unable to connect to Llama Stack: %s", e)
         response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
         raise HTTPException(**response.model_dump()) from e
 
     return ProvidersListResponse(providers=group_providers(providers))
 
 
-def group_providers(providers: ListProvidersResponse) -> dict[str, list[dict[str, Any]]]:
+def group_providers(
+    providers: ListProvidersResponse,
+) -> dict[str, list[dict[str, Any]]]:
     """Group a list of ProviderInfo objects by their API type.
 
     Args:
@@ -164,11 +166,15 @@ async def get_provider_endpoint_handler(
         provider = await client.providers.retrieve(provider_id)
         return ProviderResponse(**provider.model_dump())
 
-    except APIConnectionError as e:
-        logger.error("Unable to connect to Llama Stack: %s", e)
-        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
+    except (BadRequestError, ValueError) as e:
+        # Server mode raises BadRequestError; library mode raises ValueError.
+        logger.error("Provider not found: %s", e)
+        response = NotFoundResponse(resource="provider", resource_id=provider_id)
         raise HTTPException(**response.model_dump()) from e
 
-    except BadRequestError as e:
-        response = NotFoundResponse(resource="provider", resource_id=provider_id)
+    except ApiException as e:
+        if e.status:
+            raise
+        logger.error("Unable to connect to Llama Stack: %s", e)
+        response = ServiceUnavailableResponse(backend_name="OGX", cause=str(e))
         raise HTTPException(**response.model_dump()) from e