10xgraph-api 0.7.0__py3-none-any.whl

tenxgraph_api/src/app/utils/media/pipeline.py

@@ -0,0 +1,83 @@
+"""Document processing pipeline for the API layer.
+
+Orchestrates: receive uploaded file -> extract text via DocumentExtractor
+-> return TextBlock or DocumentBlock depending on config.
+"""
+
+from __future__ import annotations
+
+import logging
+from base64 import b64decode
+from typing import Any
+
+from tenxgraph.core.state.message_block import DocumentBlock, TextBlock
+from tenxgraph.storage.media.config import DocumentHandling
+
+from .extractor import DocumentExtractor
+
+
+logger = logging.getLogger("tenxgraph_api.media.pipeline")
+
+
+class DocumentPipeline:
+    """Document processing pipeline with extraction config.
+
+    Modes:
+        - ``EXTRACT_TEXT``: Extract text via textxtract, return TextBlock.
+        - ``PASS_RAW``/``FORWARD_RAW``: Return the original DocumentBlock untouched.
+        - ``SKIP``: Drop the document entirely (returns None).
+    """
+
+    def __init__(
+        self,
+        document_extractor: DocumentExtractor | None = None,
+        handling: DocumentHandling = DocumentHandling.EXTRACT_TEXT,
+    ):
+        self.handling = handling
+        # Lazy-init extractor only when actually needed (EXTRACT_TEXT mode)
+        self._extractor = document_extractor
+
+    @property
+    def extractor(self) -> DocumentExtractor:
+        if self._extractor is None:
+            self._extractor = DocumentExtractor()
+        return self._extractor
+
+    async def process_document(self, document_block: Any) -> Any | None:
+        """Process a DocumentBlock according to the handling policy.
+
+        Args:
+            document_block: A ``DocumentBlock`` instance.
+
+        Returns:
+            A ``TextBlock`` (if extracted), the original ``DocumentBlock``
+            (if pass-raw or extraction unsupported), or ``None`` (if skip).
+        """
+        if self.handling == DocumentHandling.SKIP:
+            return None
+
+        if self.handling == DocumentHandling.FORWARD_RAW:
+            return document_block
+
+        # EXTRACT_TEXT path
+        if not isinstance(document_block, DocumentBlock):
+            raise ValueError("Expected DocumentBlock in pipeline")
+
+        # If an excerpt already exists, return it as text
+        if document_block.excerpt and document_block.excerpt.strip():
+            return TextBlock(text=document_block.excerpt)
+
+        media = document_block.media
+
+        if media.kind == "data" and media.data_base64:
+            decoded = b64decode(media.data_base64)
+            filename = media.filename or "document.pdf"
+            extracted = await self.extractor.extract(decoded, filename, media.mime_type)
+            if extracted:
+                return TextBlock(text=extracted)
+            # Extraction returned None (unsupported type) — keep raw
+            return document_block
+
+        # For external URLs or provider file_id, we cannot extract directly.
+        # Keep as-is so API layer may resolve later.
+        return document_block

tenxgraph_api/src/app/utils/parse_output.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from pydantic import BaseModel
+
+from tenxgraph_api.src.app.core.config.settings import Settings
+
+
+def parse_state_output(settings: Settings, response: BaseModel) -> dict[str, Any]:
+    # if settings.IS_DEBUG:
+    #     return response.model_dump(exclude={"execution_meta"}, serialize_as_any=True)
+    return response.model_dump(serialize_as_any=True)
+
+
+def parse_message_output(settings: Settings, response: BaseModel) -> dict[str, Any]:
+    # if settings.IS_DEBUG:
+    #     return response.model_dump(exclude={"raw"}, serialize_as_any=True)
+    return response.model_dump(serialize_as_any=True)

tenxgraph_api/src/app/utils/response_helper.py

@@ -0,0 +1,129 @@
+"""This module contains helper functions for creating responses
+with metadata, error codes, and messages. The functions in this
+module are used to create responses with metadata, error codes,
+and messages. The success_response function is used to create
+a success response with the provided data, message, status code,
+and metadata. The error_response function is used to create an
+error response with the provided error code, message, details,
+status code, and metadata. The merge_metadata function is used
+to merge metadata with request details and a message.
+
+Examples:
+    - success_response(res, request, message, status_code, metadata)
+    - error_response(request, error_code, message, details, status_code, metadata)
+    - merge_metadata(metadata, request, message)
+
+"""
+
+from typing import Any, TypeVar
+
+from fastapi import Request
+from fastapi.responses import ORJSONResponse
+
+from .schemas import (
+    ErrorOutputSchema,
+    ErrorResponse,
+    ErrorSchemas,
+    SuccessResponse,
+)
+
+
+T = TypeVar("T")
+
+
+def merge_metadata(metadata: dict | None, request: Request, message: str = "") -> dict[Any, Any]:
+    """Merges metadata with request details and a message.
+    The metadata is updated with the request_id, timestamp,
+    and message. This function is used to create a response
+    with metadata, request details, and a message.
+
+    Args:
+        metadata: A dictionary containing metadata information. Defaults to None.
+        request: An instance of the Request class from fastapi. Defaults to None.
+        message: A string message to be included in the metadata. Defaults to "".
+
+    Returns:
+        metadata: A dictionary containing merged metadata with request details
+            and the message.
+    """
+    if metadata:
+        metadata.update(
+            {
+                "request_id": request.state.request_id,
+                "timestamp": request.state.timestamp,
+                "message": message,
+            }
+        )
+        return metadata
+
+    return {
+        "request_id": request.state.request_id,
+        "timestamp": request.state.timestamp,
+        "message": message,
+    }
+
+
+def success_response(
+    res: Any,
+    request: Request,
+    message: str = "OK",
+    status_code: int = 200,
+    metadata: dict | None = None,
+):
+    """Creates a success response with the provided data, message,
+    status code, and metadata.
+
+    Args:
+        res (T): The data to be included in the response.
+        request (Request): The FastAPI request object.
+        message (str, optional): The message associated with the response.
+            Defaults to OK.
+        status_code (int, optional): The HTTP status code of the response.
+            Defaults to 200.
+        metadata (dict, optional): Additional metadata to be merged with the response.
+            Defaults to None.
+
+    Returns:
+        ORJSONResponse: The JSON response containing the data,
+            message, and metadata.
+    """
+    response: SuccessResponse = SuccessResponse(
+        data=res, metadata=merge_metadata(metadata, request, message)
+    )
+    return ORJSONResponse(response.model_dump(), status_code=status_code)
+
+
+def error_response(
+    request: Request,
+    error_code: str,
+    message: str = "",
+    details: list[ErrorSchemas] | None = None,
+    status_code: int = 400,
+    metadata: dict | None = None,
+):
+    """Creates an error response with the provided error code, message,
+    details, status code, and metadata.
+
+    Args:
+        request (Request): The FastAPI request object. Default to None.
+        error_code (str): The code associated with the error.
+            Default to None.
+        message (str, optional): The message describing the error. Default to "".
+        details (list[ErrorSchemas], optional): Additional details about the error.
+            Default to None.
+        status_code (int, optional): The HTTP status code of the response.
+            Default is 400.
+        metadata (dict, optional): Additional metadata to be merged with the response.
+            Default is None.
+
+    Returns:
+        ORJSONResponse: The JSON response containing the error details, message,
+            and metadata.
+    """
+    error = ErrorResponse(
+        error=ErrorOutputSchema(
+            code=error_code, message=message, details=details if details else []
+        ),
+        metadata=merge_metadata(metadata, request, message),
+    )
+    return ORJSONResponse(error.model_dump(), status_code=status_code)

tenxgraph_api/src/app/utils/schemas/__init__.py

@@ -0,0 +1,11 @@
+from .output_schemas import ErrorOutputSchema, ErrorResponse, ErrorSchemas, SuccessResponse
+from .user_schemas import AuthUserSchema
+
+
+__all__ = [
+    "AuthUserSchema",
+    "ErrorOutputSchema",
+    "ErrorResponse",
+    "ErrorSchemas",
+    "SuccessResponse",
+]

tenxgraph_api/src/app/utils/schemas/output_schemas.py

@@ -0,0 +1,28 @@
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel, Field
+
+
+T = TypeVar("T")
+
+
+class ErrorSchemas(BaseModel):
+    loc: list[str | int] = Field(..., title="Location")
+    msg: str = Field(..., title="Error message")
+    type: str = Field(..., title="Error type")
+
+
+class ErrorOutputSchema(BaseModel):
+    code: str = Field(..., title="Error code")
+    message: str = Field(..., title="Error message")
+    details: list[ErrorSchemas] = Field([], title="Error details")
+
+
+class ErrorResponse(BaseModel):
+    error: ErrorOutputSchema | None = None
+    metadata: dict = Field({}, title="Metadata")
+
+
+class SuccessResponse(BaseModel, Generic[T]):
+    data: T | None = None
+    metadata: dict = Field({}, title="Metadata")

tenxgraph_api/src/app/utils/schemas/user_schemas.py

@@ -0,0 +1,13 @@
+from pydantic import BaseModel
+
+
+class AuthUserSchema(BaseModel):
+    name: str
+    role: str
+    company: int
+    uuid: str
+    user_id: str
+    email: str
+    email_verified: bool
+    firebase: dict
+    uid: str

tenxgraph_api/src/app/utils/snowflake_id_generator.py

@@ -0,0 +1,90 @@
+import os
+from importlib.util import find_spec
+
+from tenxgraph.utils.id_generator import BaseIDGenerator, IDType
+
+
+# Check if snowflakekit is available
+HAS_SNOWFLAKE = find_spec("snowflakekit") is not None
+
+
+class SnowFlakeIdGenerator(BaseIDGenerator):
+    def __init__(
+        self,
+        snowflake_epoch: int | None = None,
+        total_bits: int | None = None,
+        snowflake_time_bits: int | None = None,
+        snowflake_node_bits: int | None = None,
+        snowflake_node_id: int | None = None,
+        snowflake_worker_id: int | None = None,
+        snowflake_worker_bits: int | None = None,
+    ):
+        # IF all these are None then try to read from env
+        config = None
+        if not HAS_SNOWFLAKE:
+            raise ImportError(
+                "snowflakekit is not installed. Please install it to use SnowFlakeIdGenerator."
+            )
+
+        from snowflakekit import SnowflakeConfig, SnowflakeGenerator
+
+        if (
+            snowflake_epoch is None
+            and total_bits is None
+            and snowflake_time_bits is None
+            and snowflake_node_bits is None
+            and snowflake_node_id is None
+            and snowflake_worker_id is None
+            and snowflake_worker_bits is None
+        ):
+            snowflake_epoch = int(os.environ.get("SNOWFLAKE_EPOCH", "1723323246031"))
+            total_bits = int(os.environ.get("SNOWFLAKE_TOTAL_BITS", "64"))
+            snowflake_time_bits = int(os.environ.get("SNOWFLAKE_TIME_BITS", "39"))
+            snowflake_node_bits = int(os.environ.get("SNOWFLAKE_NODE_BITS", "7"))
+            snowflake_node_id = int(os.environ.get("SNOWFLAKE_NODE_ID", "0"))
+            snowflake_worker_id = int(os.environ.get("SNOWFLAKE_WORKER_ID", "0"))
+            snowflake_worker_bits = int(os.environ.get("SNOWFLAKE_WORKER_BITS", "5"))
+
+            config = SnowflakeConfig(
+                epoch=snowflake_epoch,
+                total_bits=total_bits,
+                time_bits=snowflake_time_bits,
+                node_bits=snowflake_node_bits,
+                node_id=snowflake_node_id,
+                worker_id=snowflake_worker_id,
+                worker_bits=snowflake_worker_bits,
+            )
+
+        elif (
+            snowflake_epoch is not None
+            and total_bits is not None
+            and snowflake_time_bits is not None
+            and snowflake_node_bits is not None
+            and snowflake_node_id is not None
+            and snowflake_worker_id is not None
+            and snowflake_worker_bits is not None
+        ):
+            config = SnowflakeConfig(
+                epoch=snowflake_epoch,
+                total_bits=total_bits,
+                time_bits=snowflake_time_bits,
+                node_bits=snowflake_node_bits,
+                node_id=snowflake_node_id,
+                worker_id=snowflake_worker_id,
+                worker_bits=snowflake_worker_bits,
+            )
+
+        # Now create generator
+        self.generator = SnowflakeGenerator(config=config)
+
+    @property
+    def id_type(self) -> IDType:
+        return IDType.BIGINT
+
+    async def generate(self) -> int:
+        """Generate a new Snowflake ID.
+
+        Returns:
+            int: A new Snowflake ID.
+        """
+        return await self.generator.generate()

tenxgraph_api/src/app/utils/swagger_helper.py

@@ -0,0 +1,160 @@
+"""
+Swagger Helper Module
+Contains helper functions to generate Swagger Responses
+"""
+
+import uuid
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel
+
+from .schemas import (
+    ErrorOutputSchema,
+    ErrorSchemas,
+)
+
+
+class _SwaggerSuccessSchemas[T](BaseModel):
+    """
+    Swagger Success Response Schema
+    """
+
+    data: T
+    metadata: dict = {
+        "message": "Success",
+        "request_id": uuid.uuid4().hex,
+        "timestamp": datetime.now().isoformat(),
+    }
+
+
+class _SwaggerSuccessPaginationSchemas[T](BaseModel):
+    data: T
+    metadata: dict = {
+        "message": "Success",
+        "request_id": uuid.uuid4().hex,
+        "timestamp": datetime.now().isoformat(),
+        "pagination": {
+            "total": 100,
+            "offset": 0,
+            "limit": 10,
+        },
+    }
+
+
+class _SwaggerError400Schemas(BaseModel):
+    metadata: dict = {
+        "message": "Failed",
+        "request_id": uuid.uuid4().hex,
+        "timestamp": datetime.now().isoformat(),
+    }
+    error: ErrorOutputSchema = ErrorOutputSchema(
+        code="BAD_REQUEST",
+        message="Invalid input, please check the input data for any errors",
+        details=[],
+    )
+
+
+class _SwaggerError404Schemas(BaseModel):
+    metadata: dict = {
+        "message": "Failed",
+        "request_id": uuid.uuid4().hex,
+        "timestamp": datetime.now().isoformat(),
+    }
+    error: ErrorOutputSchema = ErrorOutputSchema(
+        code="RESOURCE_NOT_FOUND", message="Resource not found", details=[]
+    )
+
+
+class _SwaggerError401Schemas(BaseModel):
+    metadata: dict = {
+        "message": "Failed",
+        "request_id": uuid.uuid4().hex,
+        "timestamp": datetime.now().isoformat(),
+    }
+    error: ErrorOutputSchema = ErrorOutputSchema(
+        code="AUTHENTICATION_FAILED",
+        message="Please provide valid credentials",
+        details=[],
+    )
+
+
+class _SwaggerError403Schemas(BaseModel):
+    metadata: dict = {
+        "message": "Failed",
+        "request_id": uuid.uuid4().hex,
+        "timestamp": datetime.now().isoformat(),
+    }
+    error: ErrorOutputSchema = ErrorOutputSchema(
+        code="PERMISSION_ERROR",
+        message="You don't have permission to access this resource",
+        details=[],
+    )
+
+
+class _SwaggerError426Schemas(BaseModel):
+    metadata: dict = {
+        "message": "Failed",
+        "request_id": uuid.uuid4().hex,
+        "timestamp": datetime.now().isoformat(),
+    }
+    error: ErrorOutputSchema = ErrorOutputSchema(
+        code="VALIDATION_ERROR",
+        message="Invalid input",
+        details=[
+            ErrorSchemas(
+                loc=["body", "name"],
+                msg="field required",
+                type="value_error.missing",
+            )
+        ],
+    )
+
+
+def generate_swagger_responses(
+    model: type[BaseModel],
+    show_pagination: bool = False,
+) -> dict[int | str, dict[str, Any]]:
+    """
+    Generate Swagger Responses Example, based on the provided model.
+    This will generate the following responses:
+    - 400: Invalid input
+    - 404: Resource not found
+    - 401: Authentication failed
+    - 403: Forbidden
+    - 422: Validation error
+    - 200: Success
+
+    The 200 response will be based on the provided model. If show_pagination is True,
+    then the response will have pagination metadata.
+    It will be used to generate redoc and swagger documentation.
+
+    Args:
+        model (type[BaseModel]): Base Pydantic Model class, which will be used as response.
+        show_pagination (bool, optional): Show Pagination in Swagger Response. Defaults to False.
+
+    Returns:
+        dict[int, dict[str, Any]]: Collection of Swagger Responses
+
+    """
+    return {
+        400: {"model": _SwaggerError400Schemas, "description": "Invalid input"},
+        404: {
+            "model": _SwaggerError404Schemas,
+            "description": "Resource not found",
+        },
+        401: {
+            "model": _SwaggerError401Schemas,
+            "description": "Authentication failed",
+        },
+        403: {"model": _SwaggerError403Schemas, "description": "Forbidden"},
+        422: {
+            "model": _SwaggerError426Schemas,
+            "description": "Validation error",
+        },
+        200: {
+            "model": _SwaggerSuccessPaginationSchemas[model]  # type: ignore
+            if show_pagination
+            else _SwaggerSuccessSchemas[model],  # type: ignore
+        },
+    }

tenxgraph_api/src/app/utils/telemetry_store.py

@@ -0,0 +1,123 @@
+"""In-memory telemetry store for observability.
+
+Captures per-run streaming records (one per ``StreamChunk``) keyed by thread, so
+the observability endpoint can reconstruct a trace (spans + events + token cost)
+from the events a run actually emitted. This is process-local and best-effort —
+it is not a durable telemetry backend; restart clears it.
+
+A single instance is bound as an InjectQ singleton and written to by the graph
+service during invoke/stream runs.
+"""
+
+from __future__ import annotations
+
+import threading
+from collections import deque
+from typing import Any
+
+
+# Bounds so long-lived servers don't grow unbounded.
+MAX_THREADS = 200
+MAX_RUNS_PER_THREAD = 20
+MAX_RECORDS_PER_RUN = 2000
+
+
+class RunTrace:
+    """A single graph run's captured records."""
+
+    def __init__(self, run_id: str, thread_id: str, started_at: float) -> None:
+        self.run_id = run_id
+        self.thread_id = thread_id
+        self.started_at = started_at
+        self.finished_at: float | None = None
+        self.status: str = "running"  # running | done | error | stopped
+        self.records: deque[dict[str, Any]] = deque(maxlen=MAX_RECORDS_PER_RUN)
+
+    def add(self, record: dict[str, Any]) -> None:
+        self.records.append(record)
+
+
+class TelemetryStore:
+    """Thread-safe, in-memory store of run traces keyed by thread id."""
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        # thread_id -> ordered dict-like: run_id -> RunTrace (kept insertion order)
+        self._threads: dict[str, dict[str, RunTrace]] = {}
+        self._thread_order: deque[str] = deque(maxlen=MAX_THREADS)
+
+    def start_run(self, thread_id: str, run_id: str, started_at: float) -> RunTrace:
+        thread_id = str(thread_id)
+        run_id = str(run_id)
+        with self._lock:
+            runs = self._threads.get(thread_id)
+            if runs is None:
+                runs = {}
+                self._threads[thread_id] = runs
+                if thread_id in self._thread_order:
+                    self._thread_order.remove(thread_id)
+                self._thread_order.append(thread_id)
+                # Evict oldest threads beyond the cap.
+                while len(self._threads) > MAX_THREADS and self._thread_order:
+                    oldest = self._thread_order.popleft()
+                    self._threads.pop(oldest, None)
+
+            trace = RunTrace(run_id, thread_id, started_at)
+            runs[run_id] = trace
+            # Cap runs per thread (drop oldest by insertion order).
+            while len(runs) > MAX_RUNS_PER_THREAD:
+                oldest_run = next(iter(runs))
+                runs.pop(oldest_run, None)
+            return trace
+
+    def record(self, thread_id: str, run_id: str, record: dict[str, Any]) -> None:
+        thread_id = str(thread_id)
+        run_id = str(run_id)
+        with self._lock:
+            runs = self._threads.get(thread_id)
+            if not runs:
+                return
+            trace = runs.get(run_id)
+            if trace:
+                trace.add(record)
+
+    def finish_run(
+        self,
+        thread_id: str,
+        run_id: str,
+        finished_at: float,
+        status: str = "done",
+    ) -> None:
+        thread_id = str(thread_id)
+        run_id = str(run_id)
+        with self._lock:
+            runs = self._threads.get(thread_id)
+            if not runs:
+                return
+            trace = runs.get(run_id)
+            if trace:
+                trace.finished_at = finished_at
+                trace.status = status
+
+    def get_runs(self, thread_id: str) -> list[RunTrace]:
+        """Runs for a thread, newest last (insertion order)."""
+        with self._lock:
+            runs = self._threads.get(str(thread_id))
+            return list(runs.values()) if runs else []
+
+    def get_latest_run(self, thread_id: str) -> RunTrace | None:
+        runs = self.get_runs(thread_id)
+        return runs[-1] if runs else None
+
+    def get_run(self, thread_id: str, run_id: str) -> RunTrace | None:
+        with self._lock:
+            runs = self._threads.get(str(thread_id))
+            return runs.get(str(run_id)) if runs else None
+
+    def delete_thread(self, thread_id: str) -> None:
+        """Completely clear a thread's traces from the in-memory telemetry cache."""
+        thread_id = str(thread_id)
+        with self._lock:
+            self._threads.pop(thread_id, None)
+            if thread_id in self._thread_order:
+                self._thread_order.remove(thread_id)