Update imports and mocks

Author: samdoran

scripts/generate_openapi_schema.py

@@ -7,10 +7,10 @@
 
 from fastapi.openapi.utils import get_openapi
 
-from client import AsyncLlamaStackClientHolder
+from lightspeed_stack.client import AsyncLlamaStackClientHolder
 
 # it is needed to read proper configuration in order to start the app to generate schema
-from configuration import configuration
+from lightspeed_stack.configuration import configuration
 
 cfg_file = "lightspeed-stack.yaml"
 configuration.load_configuration(cfg_file)
@@ -20,7 +20,7 @@
 
 asyncio.run(AsyncLlamaStackClientHolder().load(configuration.configuration.llama_stack))
 
-from app.main import app  # noqa: E402  pylint: disable=C0413
+from lightspeed_stack.app.main import app  # noqa: E402  pylint: disable=C0413
 
 
 def read_version_from_openapi(filename: str) -> str:

src/lightspeed_stack/__main__.py

@@ -2,4 +2,5 @@
 
 from lightspeed_stack.cli import main
 
-main()
+if __name__ == "__main__":
+    main()

src/lightspeed_stack/a2a_storage/__init__.py

@@ -8,11 +8,11 @@
 state is shared across all workers.
 """
 
-from a2a_storage.context_store import A2AContextStore
-from a2a_storage.in_memory_context_store import InMemoryA2AContextStore
-from a2a_storage.postgres_context_store import PostgresA2AContextStore
-from a2a_storage.sqlite_context_store import SQLiteA2AContextStore
-from a2a_storage.storage_factory import A2AStorageFactory
+from lightspeed_stack.a2a_storage.context_store import A2AContextStore
+from lightspeed_stack.a2a_storage.in_memory_context_store import InMemoryA2AContextStore
+from lightspeed_stack.a2a_storage.postgres_context_store import PostgresA2AContextStore
+from lightspeed_stack.a2a_storage.sqlite_context_store import SQLiteA2AContextStore
+from lightspeed_stack.a2a_storage.storage_factory import A2AStorageFactory
 
 __all__ = [
     "A2AContextStore",

src/lightspeed_stack/a2a_storage/in_memory_context_store.py

@@ -3,8 +3,8 @@
 import asyncio
 from typing import Optional
 
-from a2a_storage.context_store import A2AContextStore
-from log import get_logger
+from lightspeed_stack.a2a_storage.context_store import A2AContextStore
+from lightspeed_stack.log import get_logger
 
 logger = get_logger(__name__)
 

src/lightspeed_stack/a2a_storage/postgres_context_store.py

@@ -6,8 +6,8 @@
 from sqlalchemy.dialects.postgresql import insert as pg_insert
 from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker
 
-from a2a_storage.context_store import A2AContextStore
-from log import get_logger
+from lightspeed_stack.a2a_storage.context_store import A2AContextStore
+from lightspeed_stack.log import get_logger
 
 logger = get_logger(__name__)
 

src/lightspeed_stack/a2a_storage/sqlite_context_store.py

@@ -5,8 +5,8 @@
 from sqlalchemy import Column, MetaData, String, Table, delete, select
 from sqlalchemy.ext.asyncio import AsyncEngine, async_sessionmaker
 
-from a2a_storage.context_store import A2AContextStore
-from log import get_logger
+from lightspeed_stack.a2a_storage.context_store import A2AContextStore
+from lightspeed_stack.log import get_logger
 
 logger = get_logger(__name__)
 

src/lightspeed_stack/a2a_storage/storage_factory.py

@@ -6,12 +6,12 @@
 from a2a.server.tasks import DatabaseTaskStore, InMemoryTaskStore, TaskStore
 from sqlalchemy.ext.asyncio import AsyncEngine, create_async_engine
 
-from a2a_storage.context_store import A2AContextStore
-from a2a_storage.in_memory_context_store import InMemoryA2AContextStore
-from a2a_storage.postgres_context_store import PostgresA2AContextStore
-from a2a_storage.sqlite_context_store import SQLiteA2AContextStore
-from log import get_logger
-from models.config import A2AStateConfiguration
+from lightspeed_stack.a2a_storage.context_store import A2AContextStore
+from lightspeed_stack.a2a_storage.in_memory_context_store import InMemoryA2AContextStore
+from lightspeed_stack.a2a_storage.postgres_context_store import PostgresA2AContextStore
+from lightspeed_stack.a2a_storage.sqlite_context_store import SQLiteA2AContextStore
+from lightspeed_stack.log import get_logger
+from lightspeed_stack.models.config import A2AStateConfiguration
 
 logger = get_logger(__name__)
 

src/lightspeed_stack/app/database.py

@@ -8,10 +8,10 @@
 from sqlalchemy.engine.base import Engine
 from sqlalchemy.orm import Session, sessionmaker
 
-from configuration import configuration
-from log import get_logger
-from models.config import PostgreSQLDatabaseConfiguration, SQLiteDatabaseConfiguration
-from models.database.base import Base
+from lightspeed_stack.configuration import configuration
+from lightspeed_stack.log import get_logger
+from lightspeed_stack.models.config import PostgreSQLDatabaseConfiguration, SQLiteDatabaseConfiguration
+from lightspeed_stack.models.database.base import Base
 
 logger = get_logger(__name__)
 

src/lightspeed_stack/app/endpoints/a2a.py

@@ -49,24 +49,24 @@
 from pydantic_ai.run import AgentRunResult
 from starlette.responses import Response, StreamingResponse
 
-from a2a_storage import A2AContextStore, A2AStorageFactory
-from app.endpoints.a2a_openapi import a2a_jsonrpc_responses
-from authentication import get_auth_dependency
-from authentication.interface import AuthTuple
-from authorization.middleware import authorize
-from client import AsyncLlamaStackClientHolder
-from configuration import configuration
-from constants import MEDIA_TYPE_EVENT_STREAM
-from log import get_logger
-from models.api.requests import QueryRequest
-from models.config import Action
-from utils.agents.query import map_agent_inference_error
-from utils.conversation_compaction import apply_compaction_blocking
-from utils.mcp_headers import McpHeaders, mcp_headers_dependency
-from utils.pydantic_ai_helpers import build_agent
-from utils.responses import prepare_responses_params
-from utils.suid import normalize_conversation_id
-from version import __version__
+from lightspeed_stack.a2a_storage import A2AContextStore, A2AStorageFactory
+from lightspeed_stack.app.endpoints.a2a_openapi import a2a_jsonrpc_responses
+from lightspeed_stack.authentication import get_auth_dependency
+from lightspeed_stack.authentication.interface import AuthTuple
+from lightspeed_stack.authorization.middleware import authorize
+from lightspeed_stack.client import AsyncLlamaStackClientHolder
+from lightspeed_stack.configuration import configuration
+from lightspeed_stack.constants import MEDIA_TYPE_EVENT_STREAM
+from lightspeed_stack.log import get_logger
+from lightspeed_stack.models.api.requests import QueryRequest
+from lightspeed_stack.models.config import Action
+from lightspeed_stack.utils.agents.query import map_agent_inference_error
+from lightspeed_stack.utils.conversation_compaction import apply_compaction_blocking
+from lightspeed_stack.utils.mcp_headers import McpHeaders, mcp_headers_dependency
+from lightspeed_stack.utils.pydantic_ai_helpers import build_agent
+from lightspeed_stack.utils.responses import prepare_responses_params
+from lightspeed_stack.utils.suid import normalize_conversation_id
+from lightspeed_stack.version import __version__
 
 logger = get_logger(__name__)
 router = APIRouter(tags=["a2a"])

src/lightspeed_stack/app/endpoints/a2a_openapi.py

@@ -2,7 +2,7 @@
 
 from typing import Any, Final
 
-from constants import MEDIA_TYPE_EVENT_STREAM, MEDIA_TYPE_JSON
+from lightspeed_stack.constants import MEDIA_TYPE_EVENT_STREAM, MEDIA_TYPE_JSON
 
 # 200 may be buffered JSON-RPC (application/json) or SSE (text/event-stream).
 a2a_jsonrpc_responses: Final[dict[int | str, dict[str, Any]]] = {