refactor(schemas): clean up imports and update comments in a2a_types and schemas

src/schemas/a2a_types.py

@@ -1,6 +1,6 @@
 from datetime import datetime
 from enum import Enum
-from typing import Annotated, Any, Literal, TypeVar
+from typing import Annotated, Any, Literal
 from uuid import uuid4
 from typing_extensions import Self
 
@@ -147,7 +147,7 @@ class TaskPushNotificationConfig(BaseModel):
     pushNotificationConfig: PushNotificationConfig
 
 
-## RPC Messages
+# RPC Messages
 
 
 class JSONRPCMessage(BaseModel):
@@ -243,7 +243,7 @@ A2ARequest = TypeAdapter(
     ]
 )
 
-## Error types
+# Error types
 
 
 class JSONParseError(JSONRPCError):

src/schemas/schemas.py

@@ -1,5 +1,5 @@
-from pydantic import BaseModel, Field, validator, EmailStr, UUID4, ConfigDict
-from typing import Optional, Dict, Any, Union, List
+from pydantic import BaseModel, Field, validator, UUID4, ConfigDict
+from typing import Optional, Dict, Any, List
 from datetime import datetime
 from uuid import UUID
 import uuid

src/services/a2a_task_manager.py

@@ -1,8 +1,7 @@
-import json
 import logging
 import asyncio
 from collections.abc import AsyncIterable
-from typing import Any, Dict, Optional, Union, List
+from typing import Dict, Optional
 from uuid import UUID
 
 from sqlalchemy.orm import Session
@@ -10,19 +9,8 @@ from sqlalchemy.orm import Session
 from src.config.settings import settings
 from src.services.agent_service import (
     get_agent,
-    create_agent,
-    update_agent,
-    delete_agent,
-    get_agents_by_client,
 )
 from src.services.mcp_server_service import get_mcp_server
-from src.services.session_service import (
-    get_sessions_by_client,
-    get_sessions_by_agent,
-    get_session_by_id,
-    delete_session,
-    get_session_events,
-)
 
 from src.services.agent_runner import run_agent, run_agent_stream
 from src.services.service_providers import (
@@ -367,7 +355,7 @@ class A2ATaskManager:
             final_artifact = Artifact(parts=parts, index=0)
 
             # Update the task in the store with the final response
-            task = await self.update_store(
+            await self.update_store(
                 task_params.id,
                 TaskStatus(state=task_state, message=final_message),
                 [final_artifact],

src/services/agent_builder.py

@@ -139,9 +139,7 @@ class AgentBuilder:
                 logger.info(f"Using stored API key for agent {agent.name}")
                 api_key = decrypted_key
             else:
-                logger.error(
-                    f"Stored API key not found for agent {agent.name}"
-                )
+                logger.error(f"Stored API key not found for agent {agent.name}")
                 raise ValueError(
                     f"API key with ID {agent.api_key_id} not found or inactive"
                 )
@@ -155,7 +153,7 @@ class AgentBuilder:
                     key_id = uuid.UUID(config_api_key)
                     decrypted_key = get_decrypted_api_key(self.db, key_id)
                     if decrypted_key:
-                        logger.info(f"Config API key is a valid reference")
+                        logger.info("Config API key is a valid reference")
                         api_key = decrypted_key
                     else:
                         # Use the key directly
@@ -165,7 +163,9 @@ class AgentBuilder:
                     api_key = config_api_key
             else:
                 logger.error(f"No API key configured for agent {agent.name}")
-                raise ValueError(f"Agent {agent.name} does not have a configured API key")
+                raise ValueError(
+                    f"Agent {agent.name} does not have a configured API key"
+                )
 
         return (
             LlmAgent(

src/services/apikey_service.py

@@ -67,7 +67,7 @@ def get_api_keys_by_client(
         query = (
             db.query(ApiKey)
             .filter(ApiKey.client_id == client_id)
-            .filter(ApiKey.is_active == True)
+            .filter(ApiKey.is_active)
         )
 
         # Apply sorting

src/services/workflow_agent.py

@@ -7,7 +7,6 @@ from google.genai.types import Content, Part
 from typing import AsyncGenerator, Dict, Any, List, TypedDict
 import uuid
 
-from google.adk.runners import Runner
 from src.services.agent_service import get_agent
 
 from sqlalchemy.orm import Session
@@ -280,7 +279,7 @@ class WorkflowAgent(BaseAgent):
 
             # Prepare a more descriptive message about the conditions
             conditions_result_text = "\n".join(condition_details)
-            condition_summary = f"TRUE" if conditions_met else "FALSE"
+            condition_summary = "TRUE" if conditions_met else "FALSE"
 
             condition_content = [
                 Event(