Fix/import

.coveragerc

@@ -1,23 +1,27 @@
 [run]
 branch = True
-source = src
+source = agentflow_cli
 omit =
-    src/tests/*
+    */__init__.py
+    tests/*
+    */tests/*
     graph/*
     migrations/*
-    src/app/worker.py
-    src/app/main.py
-    src/app/loader.py
-    src/app/routers/checkpointer/*
-    src/app/routers/graph/*
-    src/app/routers/*/services/*  # services often depend on external libs; still test parts
-    src/app/core/auth/*
-    src/app/core/config/setup_logs.py
-    src/app/core/config/sentry_config.py
-    src/app/core/config/worker_middleware.py
+    scripts/*
+    venv/*
+    .venv/*
 
 [report]
 exclude_lines =
     pragma: no cover
     if TYPE_CHECKING:
     if __name__ == .__main__.
+    @abc.abstractmethod
+    @abstractmethod
+    raise NotImplementedError
+show_missing = True
+
+[paths]
+source =
+    agentflow_cli
+    */site-packages/agentflow_cli

agentflow_cli/src/app/core/auth/jwt_auth.py

@@ -1,4 +1,3 @@
-import os
 from typing import Any
 
 import jwt
@@ -7,6 +6,7 @@
 
 from agentflow_cli.src.app.core import logger
 from agentflow_cli.src.app.core.auth.base_auth import BaseAuth
+from agentflow_cli.src.app.core.config.settings import get_settings
 from agentflow_cli.src.app.core.exceptions import UserAccountError
 
 
@@ -44,13 +44,12 @@ def authenticate(
                 message="Invalid token, please login again",
                 error_code="REVOKED_TOKEN",
             )
-        jwt_secret_key = os.environ.get("JWT_SECRET_KEY", None)
-        jwt_algorithm = os.environ.get("JWT_ALGORITHM", None)
 
-        # check bearer token then remove barer prefix
+        settings = get_settings()
+        jwt_secret_key = settings.JWT_SECRET_KEY
+        jwt_algorithm = settings.JWT_ALGORITHM
+
         token = credential.credentials
-        if token.lower().startswith("bearer "):
-            token = token[7:]
 
         if jwt_secret_key is None or jwt_algorithm is None:
             raise UserAccountError(
@@ -61,8 +60,8 @@ def authenticate(
         try:
             decoded_token = jwt.decode(
                 token,
-                jwt_secret_key,  # type: ignore
-                algorithms=[jwt_algorithm],  # type: ignore
+                jwt_secret_key,
+                algorithms=[jwt_algorithm],
             )
         except jwt.ExpiredSignatureError:
             raise UserAccountError(

agentflow_cli/src/app/core/config/media_settings.py

@@ -6,6 +6,7 @@
 from enum import Enum
 from functools import lru_cache
 
+from pydantic import ConfigDict
 from pydantic_settings import BaseSettings
 
 
@@ -41,8 +42,7 @@ class MediaSettings(BaseSettings):
     MEDIA_SIGNED_URL_TTL_SECONDS: int = 3600
     MEDIA_SIGNED_URL_REFRESH_BUFFER_SECONDS: int = 60
 
-    class Config:
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 @lru_cache

agentflow_cli/src/app/core/config/settings.py

@@ -2,7 +2,7 @@
 import os
 from functools import lru_cache
 
-from pydantic import field_validator, model_validator
+from pydantic import ConfigDict, field_validator, model_validator
 from pydantic_settings import BaseSettings
 
 
@@ -102,6 +102,12 @@ class Settings(BaseSettings):
     SNOWFLAKE_NODE_BITS: int = 5
     SNOWFLAKE_WORKER_BITS: int = 8
 
+    #################################
+    ###### JWT Config ###############
+    #################################
+    JWT_SECRET_KEY: str | None = None
+    JWT_ALGORITHM: str = "HS256"
+
     @field_validator("MODE", mode="before")
     @classmethod
     def normalize_mode(cls, v: str | None) -> str:
@@ -149,8 +155,7 @@ def check_production_security(self):
 
         return self
 
-    class Config:
-        extra = "allow"
+    model_config = ConfigDict(extra="allow")
 
 
 @lru_cache

agentflow_cli/src/app/core/exceptions/handle_errors.py

@@ -1,5 +1,5 @@
 # Handle all exceptions of agentflow here
-from agentflow.exceptions import (
+from agentflow.core.exceptions import (
     GraphError,
     GraphRecursionError,
     MetricsError,

agentflow_cli/src/app/loader.py

@@ -4,9 +4,9 @@
 import os
 from pathlib import Path
 
-from agentflow.checkpointer import BaseCheckpointer
-from agentflow.graph import CompiledGraph
-from agentflow.store import BaseStore
+from agentflow.core import CompiledGraph
+from agentflow.storage.checkpointer import BaseCheckpointer
+from agentflow.storage.store import BaseStore
 from injectq import InjectQ
 
 from agentflow_cli import BaseAuth

agentflow_cli/src/app/main.py

@@ -1,6 +1,7 @@
+import logging
 import os
 
-from agentflow.graph import CompiledGraph
+from agentflow.core.graph import CompiledGraph
 from fastapi import FastAPI
 from fastapi.concurrency import asynccontextmanager
 from fastapi.responses import ORJSONResponse
@@ -21,6 +22,8 @@
 from agentflow_cli.src.app.routers import init_routes
 
 
+logger = logging.getLogger("agentflow_api")
+
 settings = get_settings()
 # redis_client = Redis(
 #     host=settings.REDIS_HOST,
@@ -36,6 +39,28 @@
 container.bind_instance(GraphConfig, graph_config)
 
 
+async def _cleanup_temp_media_cache() -> None:
+    """Run best-effort cleanup of expired temporary media cache entries."""
+    try:
+        from agentflow.storage.media.temp_cache import TemporaryMediaCache
+
+        checkpointer = container.try_get("checkpointer") or container.try_get("BaseCheckpointer")
+        media_store = container.try_get("media_store") or container.try_get("BaseMediaStore")
+
+        if checkpointer is None:
+            logger.debug("No checkpointer available, skipping temp media cache cleanup")
+            return
+
+        cache = TemporaryMediaCache()
+        cleaned = await cache.cleanup(checkpointer, media_store)
+        if cleaned:
+            logger.info("Cleaned up %d expired temporary media cache entries on startup", cleaned)
+        else:
+            logger.debug("No expired temporary media cache entries to clean up")
+    except Exception as e:
+        logger.warning("Failed to clean up temporary media cache on startup: %s", e)
+
+
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # Load the cache
@@ -46,6 +71,9 @@ async def lifespan(app: FastAPI):
         container=container,
     )
 
+    # Clean up expired temporary media cache on startup
+    await _cleanup_temp_media_cache()
+
     # load Store
     # store = load_store(graph_config.store_path)
     # injector.binder.bind(BaseStore, store)

agentflow_cli/src/app/routers/checkpointer/router.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from agentflow.state import Message
+from agentflow.core.state import Message
 from fastapi import APIRouter, Depends, HTTPException, Request, status
 from injectq.integrations import InjectAPI
 

agentflow_cli/src/app/routers/checkpointer/schemas/checkpointer_schemas.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from agentflow.state import Message
+from agentflow.core.state import Message
 from pydantic import BaseModel, Field
 
 

agentflow_cli/src/app/routers/checkpointer/services/checkpointer_service.py

@@ -1,7 +1,7 @@
 from typing import Any
 
-from agentflow.checkpointer import BaseCheckpointer
-from agentflow.state import AgentState, Message
+from agentflow.core.state import AgentState, Message
+from agentflow.storage.checkpointer import BaseCheckpointer
 from fastapi import HTTPException
 from injectq import inject, singleton
 

agentflow_cli/src/app/routers/graph/router.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from agentflow.state import StreamChunk
+from agentflow.core.state import StreamChunk
 from fastapi import APIRouter, Depends, Request
 from fastapi.logger import logger
 from fastapi.responses import StreamingResponse

agentflow_cli/src/app/routers/graph/schemas/graph_schemas.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from agentflow.state import Message
+from agentflow.core.state import Message
 from agentflow.utils import ResponseGranularity
 from pydantic import BaseModel, Field, field_validator
 

agentflow_cli/src/app/routers/graph/services/graph_service.py

@@ -3,9 +3,10 @@
 from typing import Any
 from uuid import uuid4
 
-from agentflow.checkpointer import BaseCheckpointer
-from agentflow.graph import CompiledGraph
-from agentflow.state import AgentState, Message, StreamChunk, StreamEvent
+from agentflow.core.exceptions.media_exceptions import UnsupportedMediaInputError
+from agentflow.core.graph import CompiledGraph
+from agentflow.core.state import AgentState, Message, StreamChunk, StreamEvent
+from agentflow.storage.checkpointer import BaseCheckpointer
 from agentflow.utils.thread_info import ThreadInfo
 from fastapi import HTTPException
 from injectq import InjectQ, inject, singleton
@@ -284,6 +285,9 @@ async def invoke_graph(
                 meta=meta,
             )
 
+        except UnsupportedMediaInputError as e:
+            logger.warning("Unsupported media input: %s", e.message)
+            raise HTTPException(status_code=422, detail=e.message)
         except ValueError as e:
             logger.warning(f"Graph input validation failed: {e}")
             raise HTTPException(status_code=422, detail=str(e))

agentflow_cli/src/app/routers/graph/services/multimodal_preprocessor.py

@@ -6,8 +6,8 @@
 import logging
 from typing import TYPE_CHECKING
 
-from agentflow.state import Message
-from agentflow.state.message_block import DocumentBlock, TextBlock
+from agentflow.core.state import Message
+from agentflow.core.state.message_block import DocumentBlock, TextBlock
 
 
 if TYPE_CHECKING:

agentflow_cli/src/app/routers/media/__init__.py

@@ -8,20 +8,18 @@
 import time
 from typing import Any
 
-from agentflow.checkpointer import BaseCheckpointer
-from agentflow.media.config import DocumentHandling
-from agentflow.media.storage.base import BaseMediaStore
+from agentflow.storage.checkpointer import BaseCheckpointer
+from agentflow.storage.media.config import DocumentHandling
+from agentflow.storage.media.storage.base import BaseMediaStore
 from injectq import InjectQ, inject, singleton
 
-from agentflow_cli.media._compat import DOCUMENT_PASS_RAW, ensure_document_handling_aliases
-from agentflow_cli.media.extractor import DocumentExtractor
-from agentflow_cli.media.pipeline import DocumentPipeline
 from agentflow_cli.src.app.core.config.media_settings import MediaSettings, MediaStorageType
+from agentflow_cli.src.app.utils.media.extractor import DocumentExtractor
+from agentflow_cli.src.app.utils.media.pipeline import DocumentPipeline
 
 
 logger = logging.getLogger("agentflow-cli.media")
 
-ensure_document_handling_aliases()
 
 _SIGNED_URL_NAMESPACE = "media:signed-url"
 _EXTRACTION_NAMESPACE = "media:extraction"
@@ -41,17 +39,17 @@ def _create_media_store(settings: MediaSettings) -> BaseMediaStore:
     stype = settings.MEDIA_STORAGE_TYPE
 
     if stype == MediaStorageType.MEMORY:
-        from agentflow.media.storage.memory_store import InMemoryMediaStore
+        from agentflow.storage.media.storage.memory_store import InMemoryMediaStore
 
         return InMemoryMediaStore()
 
     if stype == MediaStorageType.LOCAL:
-        from agentflow.media.storage.local_store import LocalFileMediaStore
+        from agentflow.storage.media.storage.local_store import LocalFileMediaStore
 
         return LocalFileMediaStore(base_dir=settings.MEDIA_STORAGE_PATH)
 
     if stype == MediaStorageType.CLOUD:
-        from agentflow.media.storage.cloud_store import CloudMediaStore
+        from agentflow.storage.media.storage.cloud_store import CloudMediaStore
         from cloud_storage_manager import (
             AwsConfig,
             CloudStorageFactory,
@@ -117,7 +115,7 @@ def _create_media_store(settings: MediaSettings) -> BaseMediaStore:
 def _create_document_pipeline(settings: MediaSettings) -> DocumentPipeline:
     handling_map = {
         "extract_text": DocumentHandling.EXTRACT_TEXT,
-        "pass_raw": DOCUMENT_PASS_RAW,
+        "pass_raw": DocumentHandling.FORWARD_RAW,
         "skip": DocumentHandling.SKIP,
     }
     handling = handling_map.get(settings.DOCUMENT_HANDLING, DocumentHandling.EXTRACT_TEXT)

agentflow_cli/src/app/routers/store/schemas/store_schemas.py

@@ -4,8 +4,8 @@
 
 from typing import Any
 
-from agentflow.state import Message
-from agentflow.store.store_schema import (
+from agentflow.core.state import Message
+from agentflow.storage.store.store_schema import (
     DistanceMetric,
     MemoryRecord,
     MemorySearchResult,

agentflow_cli/src/app/routers/store/services/store_service.py

@@ -2,8 +2,8 @@
 
 from typing import Any
 
-from agentflow.state import Message
-from agentflow.store import BaseStore
+from agentflow.core.state import Message
+from agentflow.storage.store import BaseStore
 from fastapi import HTTPException
 from injectq import inject, singleton