Fix unittest for conversations API addition

Author: luis5tb

pyproject.toml

@@ -61,6 +61,12 @@ exclude = [
   # service/ols/src/auth/k8s.py and currently has 58 Pyright issues. It
   # might need to be rewritten down the line.
   "src/authentication/k8s.py",
+  # Agent API v1 endpoints - deprecated API but still supported
+  # Type errors due to llama-stack-client not exposing Agent API types
+  "src/app/endpoints/conversations.py",
+  "src/app/endpoints/query.py",
+  "src/app/endpoints/streaming_query.py",
+  "src/utils/endpoints.py",
 ]
 extraPaths = ["./src"]
 

src/app/endpoints/conversations_v3.py

@@ -4,7 +4,7 @@
 from typing import Any
 
 from fastapi import APIRouter, Depends, HTTPException, Request, status
-from llama_stack_client import APIConnectionError, NotFoundError
+from llama_stack_client import APIConnectionError, NOT_GIVEN, NotFoundError
 
 from app.database import get_session
 from authentication import get_auth_dependency
@@ -15,13 +15,13 @@
 from models.database.conversations import UserConversation
 from models.requests import ConversationUpdateRequest
 from models.responses import (
-    AccessDeniedResponse,
     BadRequestResponse,
     ConversationDeleteResponse,
     ConversationDetails,
     ConversationResponse,
     ConversationsListResponse,
     ConversationUpdateResponse,
+    ForbiddenResponse,
     NotFoundResponse,
     ServiceUnavailableResponse,
     UnauthorizedResponse,
@@ -55,7 +55,7 @@
         "description": "Unauthorized: Invalid or missing Bearer token",
     },
     403: {
-        "model": AccessDeniedResponse,
+        "model": ForbiddenResponse,
         "description": "Client does not have permission to access conversation",
     },
     404: {
@@ -82,7 +82,7 @@
         "description": "Unauthorized: Invalid or missing Bearer token",
     },
     403: {
-        "model": AccessDeniedResponse,
+        "model": ForbiddenResponse,
         "description": "Client does not have permission to access conversation",
     },
     404: {
@@ -124,7 +124,7 @@
         "description": "Unauthorized: Invalid or missing Bearer token",
     },
     403: {
-        "model": AccessDeniedResponse,
+        "model": ForbiddenResponse,
         "description": "Client does not have permission to access conversation",
     },
     404: {
@@ -283,7 +283,7 @@ async def get_conversation_endpoint_handler(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail=BadRequestResponse(
                 resource="conversation", resource_id=conversation_id
-            ).dump_detail(),
+            ).model_dump(),
         )
 
     # Normalize the conversation ID for database operations (strip conv_ prefix if present)
@@ -309,12 +309,11 @@ async def get_conversation_endpoint_handler(
         )
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
-            detail=AccessDeniedResponse(
-                user_id=user_id,
-                resource="conversation",
-                resource_id=normalized_conv_id,
+            detail=ForbiddenResponse.conversation(
                 action="read",
-            ).dump_detail(),
+                resource_id=normalized_conv_id,
+                user_id=user_id,
+            ).model_dump(),
         )
 
     # If reached this, user is authorized to retrieve this conversation
@@ -342,8 +341,6 @@ async def get_conversation_endpoint_handler(
         )
 
         # Use Conversations API to retrieve conversation items
-        from llama_stack_client import NOT_GIVEN
-
         conversation_items_response = await client.conversations.items.list(
             conversation_id=llama_stack_conv_id,
             after=NOT_GIVEN,  # No pagination cursor
@@ -384,7 +381,7 @@ async def get_conversation_endpoint_handler(
             status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
             detail=ServiceUnavailableResponse(
                 backend_name="Llama Stack", cause=str(e)
-            ).dump_detail(),
+            ).model_dump(),
         ) from e
 
     except NotFoundError as e:
@@ -393,7 +390,7 @@ async def get_conversation_endpoint_handler(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=NotFoundResponse(
                 resource="conversation", resource_id=normalized_conv_id
-            ).dump_detail(),
+            ).model_dump(),
         ) from e
 
     except HTTPException:
@@ -402,11 +399,15 @@ async def get_conversation_endpoint_handler(
     except Exception as e:
         # Handle case where conversation doesn't exist or other errors
         logger.exception("Error retrieving conversation %s: %s", normalized_conv_id, e)
+        error_msg = (
+            f"Unknown error while getting conversation {normalized_conv_id} : "
+            f"{str(e)}"
+        )
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail={
                 "response": "Unknown error",
-                "cause": f"Unknown error while getting conversation {normalized_conv_id} : {str(e)}",
+                "cause": error_msg,
             },
         ) from e
 
@@ -444,7 +445,7 @@ async def delete_conversation_endpoint_handler(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail=BadRequestResponse(
                 resource="conversation", resource_id=conversation_id
-            ).dump_detail(),
+            ).model_dump(),
         )
 
     # Normalize the conversation ID for database operations (strip conv_ prefix if present)
@@ -465,12 +466,11 @@ async def delete_conversation_endpoint_handler(
         )
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
-            detail=AccessDeniedResponse(
-                user_id=user_id,
-                resource="conversation",
-                resource_id=normalized_conv_id,
+            detail=ForbiddenResponse.conversation(
                 action="delete",
-            ).dump_detail(),
+                resource_id=normalized_conv_id,
+                user_id=user_id,
+            ).model_dump(),
         )
 
     # If reached this, user is authorized to delete this conversation
@@ -480,7 +480,7 @@ async def delete_conversation_endpoint_handler(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=NotFoundResponse(
                 resource="conversation", resource_id=normalized_conv_id
-            ).dump_detail(),
+            ).model_dump(),
         )
 
     logger.info("Deleting conversation %s using Conversations API", normalized_conv_id)
@@ -502,16 +502,15 @@ async def delete_conversation_endpoint_handler(
 
         return ConversationDeleteResponse(
             conversation_id=normalized_conv_id,
-            success=True,
-            response="Conversation deleted successfully",
+            deleted=True,
         )
 
     except APIConnectionError as e:
         raise HTTPException(
             status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
             detail=ServiceUnavailableResponse(
                 backend_name="Llama Stack", cause=str(e)
-            ).dump_detail(),
+            ).model_dump(),
         ) from e
 
     except NotFoundError:
@@ -524,8 +523,7 @@ async def delete_conversation_endpoint_handler(
 
         return ConversationDeleteResponse(
             conversation_id=normalized_conv_id,
-            success=True,
-            response="Conversation deleted successfully",
+            deleted=True,
         )
 
     except HTTPException:
@@ -534,11 +532,15 @@ async def delete_conversation_endpoint_handler(
     except Exception as e:
         # Handle case where conversation doesn't exist or other errors
         logger.exception("Error deleting conversation %s: %s", normalized_conv_id, e)
+        error_msg = (
+            f"Unknown error while deleting conversation {normalized_conv_id} : "
+            f"{str(e)}"
+        )
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail={
                 "response": "Unknown error",
-                "cause": f"Unknown error while deleting conversation {normalized_conv_id} : {str(e)}",
+                "cause": error_msg,
             },
         ) from e
 
@@ -574,7 +576,7 @@ async def update_conversation_endpoint_handler(
             status_code=status.HTTP_400_BAD_REQUEST,
             detail=BadRequestResponse(
                 resource="conversation", resource_id=conversation_id
-            ).dump_detail(),
+            ).model_dump(),
         )
 
     # Normalize the conversation ID for database operations (strip conv_ prefix if present)
@@ -595,12 +597,11 @@ async def update_conversation_endpoint_handler(
         )
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
-            detail=AccessDeniedResponse(
-                user_id=user_id,
-                resource="conversation",
-                resource_id=normalized_conv_id,
+            detail=ForbiddenResponse.conversation(
                 action="update",
-            ).dump_detail(),
+                resource_id=normalized_conv_id,
+                user_id=user_id,
+            ).model_dump(),
         )
 
     # If reached this, user is authorized to update this conversation
@@ -610,7 +611,7 @@ async def update_conversation_endpoint_handler(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=NotFoundResponse(
                 resource="conversation", resource_id=normalized_conv_id
-            ).dump_detail(),
+            ).model_dump(),
         )
 
     logger.info(
@@ -629,7 +630,7 @@ async def update_conversation_endpoint_handler(
         metadata = {"topic_summary": update_request.topic_summary}
 
         # Use Conversations API to update the conversation metadata
-        await client.conversations.update_conversation(
+        await client.conversations.update(
             conversation_id=llama_stack_conv_id,
             metadata=metadata,
         )
@@ -663,15 +664,15 @@ async def update_conversation_endpoint_handler(
             status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
             detail=ServiceUnavailableResponse(
                 backend_name="Llama Stack", cause=str(e)
-            ).dump_detail(),
+            ).model_dump(),
         ) from e
 
     except NotFoundError as e:
         raise HTTPException(
             status_code=status.HTTP_404_NOT_FOUND,
             detail=NotFoundResponse(
                 resource="conversation", resource_id=normalized_conv_id
-            ).dump_detail(),
+            ).model_dump(),
         ) from e
 
     except HTTPException:
@@ -680,10 +681,14 @@ async def update_conversation_endpoint_handler(
     except Exception as e:
         # Handle case where conversation doesn't exist or other errors
         logger.exception("Error updating conversation %s: %s", normalized_conv_id, e)
+        error_msg = (
+            f"Unknown error while updating conversation {normalized_conv_id} : "
+            f"{str(e)}"
+        )
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail={
                 "response": "Unknown error",
-                "cause": f"Unknown error while updating conversation {normalized_conv_id} : {str(e)}",
+                "cause": error_msg,
             },
         ) from e

src/utils/endpoints.py

@@ -306,16 +306,19 @@ async def get_agent(
             existing_agent_id = agent_response.agent_id
 
     logger.debug("Creating new agent")
+    # pylint: disable=unexpected-keyword-arg,no-member
     agent = AsyncAgent(
         client,  # type: ignore[arg-type]
         model=model_id,
         instructions=system_prompt,
+        # type: ignore[call-arg]
         input_shields=available_input_shields if available_input_shields else [],
+        # type: ignore[call-arg]
         output_shields=available_output_shields if available_output_shields else [],
         tool_parser=None if no_tools else GraniteToolParser.get_parser(model_id),
-        enable_session_persistence=True,
+        enable_session_persistence=True,  # type: ignore[call-arg]
     )
-    await agent.initialize()
+    await agent.initialize()  # type: ignore[attr-defined]
 
     if existing_agent_id and conversation_id:
         logger.debug("Existing conversation ID: %s", conversation_id)
@@ -335,11 +338,12 @@ async def get_agent(
             raise HTTPException(**response.model_dump()) from e
     else:
         conversation_id = agent.agent_id
+        # pylint: enable=unexpected-keyword-arg,no-member
         logger.debug("New conversation ID: %s", conversation_id)
         session_id = await agent.create_session(get_suid())
         logger.debug("New session ID: %s", session_id)
 
-    return agent, conversation_id, session_id
+    return agent, conversation_id, session_id  # type: ignore[return-value]
 
 
 async def get_temp_agent(
@@ -360,16 +364,19 @@ async def get_temp_agent(
         tuple[AsyncAgent, str]: A tuple containing the agent and session_id.
     """
     logger.debug("Creating temporary agent")
+    # pylint: disable=unexpected-keyword-arg,no-member
     agent = AsyncAgent(
         client,  # type: ignore[arg-type]
         model=model_id,
         instructions=system_prompt,
-        enable_session_persistence=False,  # Temporary agent doesn't need persistence
+        # type: ignore[call-arg]  # Temporary agent doesn't need persistence
+        enable_session_persistence=False,
     )
-    await agent.initialize()
+    await agent.initialize()  # type: ignore[attr-defined]
 
     # Generate new IDs for the temporary agent
     conversation_id = agent.agent_id
+    # pylint: enable=unexpected-keyword-arg,no-member
     session_id = await agent.create_session(get_suid())
 
     return agent, session_id, conversation_id

src/utils/types.py

@@ -3,8 +3,10 @@
 from typing import Any, Optional
 import json
 from llama_stack_client.lib.agents.tool_parser import ToolParser
-from llama_stack_client.types.shared.completion_message import CompletionMessage
-from llama_stack_client.types.shared.tool_call import ToolCall
+from llama_stack_client.lib.agents.types import (
+    CompletionMessage as AgentCompletionMessage,
+    ToolCall as AgentToolCall,
+)
 from llama_stack_client.types.shared.interleaved_content_item import (
     TextContentItem,
     ImageContentItem,
@@ -58,16 +60,18 @@ def __call__(cls, *args, **kwargs):  # type: ignore
 class GraniteToolParser(ToolParser):
     """Workaround for 'tool_calls' with granite models."""
 
-    def get_tool_calls(self, output_message: CompletionMessage) -> list[ToolCall]:
+    def get_tool_calls(
+        self, output_message: AgentCompletionMessage
+    ) -> list[AgentToolCall]:
         """
         Return the `tool_calls` list from a CompletionMessage, or an empty list if none are present.
 
         Parameters:
-            output_message (CompletionMessage | None): Completion
+            output_message (AgentCompletionMessage | None): Completion
             message potentially containing `tool_calls`.
 
         Returns:
-            list[ToolCall]: The list of tool call entries
+            list[AgentToolCall]: The list of tool call entries
             extracted from `output_message`, or an empty list.
         """
         if output_message and output_message.tool_calls: