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

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

@@ -0,0 +1,200 @@
+"""Store API schemas."""
+
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel, Field
+from tenxgraph.core.state import Message
+from tenxgraph.storage.store.store_schema import (
+    DistanceMetric,
+    MemoryRecord,
+    MemorySearchResult,
+    MemoryType,
+    RetrievalStrategy,
+)
+
+
+# Upper bounds so one request cannot make the store (and the embedding / vector backend)
+# return or rank an unbounded number of rows.
+MAX_SEARCH_LIMIT = 100
+MAX_LIST_LIMIT = 1000
+MAX_SEARCH_TOKENS = 16_000
+
+
+class BaseConfigSchema(BaseModel):
+    """Base schema containing configuration overrides and store options."""
+
+    config: dict[str, Any] | None = Field(
+        default_factory=dict,
+        description="Configuration values forwarded to the store backend.",
+    )
+    options: dict[str, Any] | None = Field(
+        default=None,
+        description="Extra keyword arguments to forward to the store backend.",
+    )
+
+
+class StoreMemorySchema(BaseConfigSchema):
+    """Schema for storing a memory item."""
+
+    content: str | Message = Field(..., description="Memory content or structured message.")
+    memory_type: MemoryType = Field(
+        default=MemoryType.EPISODIC,
+        description="Memory classification used by the backend store.",
+    )
+    category: str = Field(default="general", description="Category label for the memory.")
+    metadata: dict[str, Any] | None = Field(
+        default=None,
+        description="Arbitrary metadata associated with the memory.",
+    )
+
+
+class SearchMemorySchema(BaseConfigSchema):
+    """Schema for searching memories."""
+
+    query: str = Field(..., description="Textual query used for memory retrieval.")
+    memory_type: MemoryType | None = Field(
+        default=None,
+        description="Optional memory type filter.",
+    )
+    category: str | None = Field(
+        default=None,
+        description="Optional category filter.",
+    )
+    limit: int = Field(
+        default=10,
+        gt=0,
+        le=MAX_SEARCH_LIMIT,
+        description="Maximum number of results to return.",
+    )
+    score_threshold: float | None = Field(
+        default=None,
+        description="Minimum similarity score required for results.",
+    )
+    filters: dict[str, Any] | None = Field(
+        default=None,
+        description="Additional store-specific filters.",
+    )
+    retrieval_strategy: RetrievalStrategy = Field(
+        default=RetrievalStrategy.SIMILARITY,
+        description="Retrieval strategy used by the backend store.",
+    )
+    distance_metric: DistanceMetric = Field(
+        default=DistanceMetric.COSINE,
+        description="Distance metric applied during similarity search.",
+    )
+    max_tokens: int = Field(
+        default=4000,
+        gt=0,
+        le=MAX_SEARCH_TOKENS,
+        description="Maximum tokens used for truncation in similarity search.",
+    )
+
+
+class UpdateMemorySchema(BaseConfigSchema):
+    """Schema for updating a memory."""
+
+    content: str | Message = Field(..., description="Updated memory content or message.")
+    metadata: dict[str, Any] | None = Field(
+        default=None,
+        description="Updated metadata for the memory.",
+    )
+
+
+class DeleteMemorySchema(BaseConfigSchema):
+    """Schema for deleting a memory."""
+
+
+class GetMemorySchema(BaseConfigSchema):
+    """Schema for retrieving a single memory."""
+
+
+class ListMemoriesSchema(BaseConfigSchema):
+    """Schema for listing memories."""
+
+    limit: int = Field(
+        default=100,
+        gt=0,
+        le=MAX_LIST_LIMIT,
+        description="Maximum number of memories to return.",
+    )
+
+
+class ForgetMemorySchema(BaseConfigSchema):
+    """Schema for forgetting memories based on filters."""
+
+    memory_type: MemoryType | None = Field(
+        default=None,
+        description="Optional memory type to target for deletion.",
+    )
+    category: str | None = Field(
+        default=None,
+        description="Optional category to target for deletion.",
+    )
+    filters: dict[str, Any] | None = Field(
+        default=None,
+        description="Additional filters to control which memories are forgotten.",
+    )
+
+
+class MemoryCreateResponseSchema(BaseModel):
+    """Response schema for create memory operations."""
+
+    memory_id: str = Field(..., description="Identifier of the stored memory.")
+
+
+class MemoryItemResponseSchema(BaseModel):
+    """Response schema for single memory retrieval."""
+
+    memory: MemorySearchResult | None = Field(
+        default=None,
+        description="Memory retrieved from the store, if available.",
+    )
+
+
+class MemoryListResponseSchema(BaseModel):
+    """Response schema for listing memories."""
+
+    memories: list[MemorySearchResult] = Field(
+        default_factory=list,
+        description="Collection of memories returned from the store.",
+    )
+
+
+class MemorySearchResponseSchema(BaseModel):
+    """Response schema for search operations."""
+
+    results: list[MemorySearchResult] = Field(
+        default_factory=list,
+        description="Search results ranked by relevance.",
+    )
+
+
+class MemoryOperationResponseSchema(BaseModel):
+    """Generic response schema for mutation operations."""
+
+    success: bool = Field(..., description="Whether the store operation succeeded.")
+    data: Any | None = Field(default=None, description="Optional payload returned by the store.")
+
+
+__all__ = [
+    "BaseConfigSchema",
+    "DeleteMemorySchema",
+    "DistanceMetric",
+    "ForgetMemorySchema",
+    "GetMemorySchema",
+    "ListMemoriesSchema",
+    "MemoryCreateResponseSchema",
+    "MemoryItemResponseSchema",
+    "MemoryListResponseSchema",
+    "MemoryOperationResponseSchema",
+    "MemoryRecord",
+    "MemorySearchResponseSchema",
+    "MemorySearchResult",
+    "MemoryType",
+    "RetrievalStrategy",
+    "SearchMemorySchema",
+    "StoreMemorySchema",
+    "UpdateMemorySchema",
+]

tenxgraph_api/src/app/routers/store/services/__init__.py

@@ -0,0 +1,4 @@
+from .store_service import StoreService
+
+
+__all__ = ["StoreService"]

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

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+from typing import Any
+
+from fastapi import HTTPException
+from injectq import inject, singleton
+from tenxgraph.core.state import Message
+from tenxgraph.storage.store import BaseStore
+
+from tenxgraph_api.src.app.core import logger
+from tenxgraph_api.src.app.core.auth.request_config import client_config
+from tenxgraph_api.src.app.routers.store.schemas.store_schemas import (
+    ForgetMemorySchema,
+    MemoryCreateResponseSchema,
+    MemoryItemResponseSchema,
+    MemoryListResponseSchema,
+    MemoryOperationResponseSchema,
+    MemorySearchResponseSchema,
+    SearchMemorySchema,
+    StoreMemorySchema,
+    UpdateMemorySchema,
+)
+
+
+class StoreUnavailableError(HTTPException, ValueError):
+    """Raised when the store service has not been configured."""
+
+    def __init__(self) -> None:
+        detail = "Store is not configured"
+        HTTPException.__init__(self, status_code=503, detail=detail)
+        ValueError.__init__(self, detail)
+
+
+# Memory fields the store owns. Client metadata is kept, but these keys are dropped from it:
+# user_id and thread_id decide who sees a memory, and memory_id is the id Mem0 reports back.
+SYSTEM_METADATA_KEYS = frozenset(
+    {"content", "user_id", "thread_id", "memory_type", "category", "timestamp", "memory_id"}
+)
+
+# Store keyword options only the server sets. ``astore`` upserts, so a client-chosen
+# memory_id could replace an existing memory; the store always generates the id instead.
+SERVER_OPTION_KEYS = frozenset({"memory_id"})
+
+# Config keys that route storage. A client-chosen collection could read another application's
+# collection or create collections without limit.
+SERVER_CONFIG_KEYS = frozenset({"collection"})
+
+
+def _client_metadata(metadata: dict[str, Any] | None) -> dict[str, Any] | None:
+    if metadata is None:
+        return None
+    return {k: v for k, v in metadata.items() if k not in SYSTEM_METADATA_KEYS}
+
+
+def _client_options(options: dict[str, Any] | None) -> dict[str, Any]:
+    return {k: v for k, v in (options or {}).items() if k not in SERVER_OPTION_KEYS}
+
+
+@singleton
+class StoreService:
+    """Service layer wrapping interactions with the configured BaseStore."""
+
+    @inject
+    def __init__(self, store: BaseStore | None):
+        self.store = store
+
+    def _get_store(self) -> BaseStore:
+        if not self.store:
+            raise StoreUnavailableError()
+        return self.store
+
+    def _config(self, config: dict[str, Any] | None, user: dict[str, Any]) -> dict[str, Any]:
+        # Client keys pass through; server-owned keys (authz, user, internals) never do.
+        # The trusted user object carries the authz policy the store enforces.
+        cfg: dict[str, Any] = {
+            k: v for k, v in client_config(config).items() if k not in SERVER_CONFIG_KEYS
+        }
+        cfg["user"] = user
+        cfg["user_id"] = user.get("user_id", "anonymous")
+        return cfg
+
+    async def store_memory(
+        self,
+        payload: StoreMemorySchema,
+        user: dict[str, Any],
+    ) -> MemoryCreateResponseSchema:
+        store = self._get_store()
+        cfg = self._config(payload.config, user)
+        options = _client_options(payload.options)
+
+        if isinstance(payload.content, Message):
+            content: str | Message = payload.content
+        else:
+            content = payload.content
+
+        memory_id = await store.astore(
+            cfg,
+            content,
+            memory_type=payload.memory_type,
+            category=payload.category,
+            metadata=_client_metadata(payload.metadata),
+            **options,
+        )
+        logger.debug("Stored memory with id %s", memory_id)
+        return MemoryCreateResponseSchema(memory_id=memory_id)
+
+    async def search_memories(
+        self,
+        payload: SearchMemorySchema,
+        user: dict[str, Any],
+    ) -> MemorySearchResponseSchema:
+        store = self._get_store()
+        cfg = self._config(payload.config, user)
+        options = _client_options(payload.options)
+
+        results = await store.asearch(
+            cfg,
+            payload.query,
+            memory_type=payload.memory_type,
+            category=payload.category,
+            limit=payload.limit,
+            score_threshold=payload.score_threshold,
+            filters=payload.filters,
+            retrieval_strategy=payload.retrieval_strategy,
+            distance_metric=payload.distance_metric,
+            max_tokens=payload.max_tokens,
+            **options,
+        )
+        return MemorySearchResponseSchema(results=results)
+
+    async def get_memory(
+        self,
+        memory_id: str,
+        config: dict[str, Any] | None,
+        user: dict[str, Any],
+        options: dict[str, Any] | None = None,
+    ) -> MemoryItemResponseSchema:
+        store = self._get_store()
+        cfg = self._config(config, user)
+        result = await store.aget(cfg, memory_id, **_client_options(options))
+        return MemoryItemResponseSchema(memory=result)
+
+    async def list_memories(
+        self,
+        config: dict[str, Any] | None,
+        user: dict[str, Any],
+        limit: int = 100,
+        options: dict[str, Any] | None = None,
+    ) -> MemoryListResponseSchema:
+        store = self._get_store()
+        cfg = self._config(config, user)
+        memories = await store.aget_all(cfg, limit=limit, **_client_options(options))
+        return MemoryListResponseSchema(memories=memories)
+
+    async def update_memory(
+        self,
+        memory_id: str,
+        payload: UpdateMemorySchema,
+        user: dict[str, Any],
+    ) -> MemoryOperationResponseSchema:
+        store = self._get_store()
+        cfg = self._config(payload.config, user)
+        options = _client_options(payload.options)
+
+        result = await store.aupdate(
+            cfg,
+            memory_id,
+            payload.content,
+            metadata=_client_metadata(payload.metadata),
+            **options,
+        )
+        return MemoryOperationResponseSchema(success=True, data=result)
+
+    async def delete_memory(
+        self,
+        memory_id: str,
+        config: dict[str, Any] | None,
+        user: dict[str, Any],
+        options: dict[str, Any] | None = None,
+    ) -> MemoryOperationResponseSchema:
+        store = self._get_store()
+        cfg = self._config(config, user)
+        result = await store.adelete(cfg, memory_id, **_client_options(options))
+        return MemoryOperationResponseSchema(success=True, data=result)
+
+    async def forget_memory(
+        self,
+        payload: ForgetMemorySchema,
+        user: dict[str, Any],
+    ) -> MemoryOperationResponseSchema:
+        store = self._get_store()
+        cfg = self._config(payload.config, user)
+        options = _client_options(payload.options)
+        forget_kwargs: dict[str, Any] = {
+            "memory_type": payload.memory_type,
+            "category": payload.category,
+            "filters": payload.filters,
+        }
+        # Remove None values before forwarding to the store
+        forget_kwargs = {k: v for k, v in forget_kwargs.items() if v is not None}
+        forget_kwargs.update(options)
+        result = await store.aforget_memory(cfg, **forget_kwargs)
+        return MemoryOperationResponseSchema(success=True, data=result)

tenxgraph_api/src/app/tasks/user_tasks.py

@@ -0,0 +1,38 @@
+# from uuid import UUID
+
+# from fastapi_injector import InjectedTaskiq
+
+# from src.app.routers.auth.services import UserService
+# from src.app.worker import broker
+
+
+# @broker.task(task_name="add_task_math")
+# def add_task_math(x: int, y: int):
+#     """
+#     Adds two integers together.
+
+#     Args:
+#         x (int): The first integer.
+#         y (int): The second integer.
+
+#     Returns:
+#         int: The sum of the two integers.
+#     """
+#     return x + y
+
+
+# @broker.task(task_name="user_service_post_processing")
+# async def post_processing_user(user_id: int, service: UserService = InjectedTaskiq(UserService)):
+#     """
+#     Post-processes a user by fetching user details using the provided user service.
+
+#     Args:
+#         user_id (int): The ID of the user to be processed.
+#         service (UserService, optional): An instance of UserService to fetch user details.
+#             Defaults to an injected instance of UserService.
+
+#     Returns:
+#         dict: A dictionary representation of the user details.
+#     """
+#     user = await service.get_user(UUID("2b1e2ae4-ff8a-48f3-8729-be0d754432e9"))
+#     return user.dict()

tenxgraph_api/src/app/utils/__init__.py

@@ -0,0 +1,12 @@
+from .response_helper import error_response, success_response
+from .swagger_helper import generate_swagger_responses
+from .thread_name_generator import DefaultThreadNameGenerator, ThreadNameGenerator
+
+
+__all__ = [
+    "DefaultThreadNameGenerator",
+    "ThreadNameGenerator",
+    "error_response",
+    "generate_swagger_responses",
+    "success_response",
+]

tenxgraph_api/src/app/utils/callable_helper.py

@@ -0,0 +1,26 @@
+import asyncio
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+
+def _is_async_callable(func: Callable[..., Any]) -> bool:
+    """Detect if a callable is a coroutine function (async def)."""
+    return asyncio.iscoroutinefunction(func)
+
+
+async def call_sync_or_async(func: Callable[..., Any], *args, **kwargs) -> Any:
+    """
+    Call a function that may be sync or async, returning its result.
+    If sync, runs in a thread pool to avoid blocking the event loop.
+    Thread pool is shared for efficiency.
+    """
+    if _is_async_callable(func):
+        return await func(*args, **kwargs)
+
+    # Call sync function in a thread pool
+    result = await asyncio.to_thread(func, *args, **kwargs)
+    # If the result is awaitable, await it
+    if inspect.isawaitable(result):
+        return await result
+    return result

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

@@ -0,0 +1,15 @@
+"""Media processing module for 10xgraph-api.
+
+Document extraction lives here (not in 10xGraph core) because
+the core library stays lightweight — SDK users extract text themselves.
+The API platform auto-extracts using textxtract.
+"""
+
+from .extractor import DocumentExtractor
+from .pipeline import DocumentPipeline
+
+
+__all__ = [
+    "DocumentExtractor",
+    "DocumentPipeline",
+]

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

@@ -0,0 +1,172 @@
+"""Document extraction service using textxtract.
+
+Wraps ``AsyncTextExtractor`` from the textxtract library.
+This module lives in 10xgraph-api, NOT in the core 10xGraph library,
+because extraction is an API-platform concern.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import io
+import logging
+import zipfile
+from pathlib import PurePosixPath
+from typing import Any
+
+
+logger = logging.getLogger("tenxgraph_api.media.extractor")
+
+try:
+    from textxtract import AsyncTextExtractor
+    from textxtract.core.exceptions import ExtractionError, FileTypeNotSupportedError
+except ImportError:  # pragma: no cover
+    AsyncTextExtractor = None  # type: ignore[assignment]
+
+
+# The handler textxtract runs is chosen from the file extension. Uploads never pass their
+# own file name through: the declared MIME type picks one of these extensions, and a file
+# with any other type (archives included) is not extracted. That keeps a file named
+# ``a.zip`` out of textxtract's ZIP handler, which unpacks nested archives with no overall
+# size budget.
+HANDLER_EXTENSIONS: dict[str, str] = {
+    "application/pdf": ".pdf",
+    "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
+    "application/msword": ".doc",
+    "text/html": ".html",
+    "text/xml": ".xml",
+    "application/xml": ".xml",
+    "text/markdown": ".md",
+    "text/csv": ".csv",
+    "application/json": ".json",
+    "text/plain": ".txt",
+}
+_EXTRACTABLE_EXTENSIONS = frozenset(HANDLER_EXTENSIONS.values())
+
+# DOCX files are ZIP containers, so they are inspected before parsing: a small upload can
+# declare gigabytes of XML.
+MAX_ARCHIVE_UNCOMPRESSED_BYTES = 50 * 1024 * 1024
+MAX_COMPRESSION_RATIO = 100
+MAX_ARCHIVE_ENTRIES = 10_000
+
+EXTRACTION_TIMEOUT_SECONDS = 30.0
+MAX_EXTRACTED_CHARS = 1_000_000
+
+
+def _handler_extension(filename: str | None, mime_type: str | None) -> str | None:
+    """The extension (and so the handler) to extract with, or None to skip extraction."""
+    if mime_type:
+        return HANDLER_EXTENSIONS.get(mime_type.split(";", 1)[0].strip().lower())
+    suffix = PurePosixPath(filename or "").suffix.lower()
+    return suffix if suffix in _EXTRACTABLE_EXTENSIONS else None
+
+
+def _check_zip_container(data: bytes) -> None:
+    """Reject a ZIP-based document whose contents would expand far beyond its size."""
+    try:
+        with zipfile.ZipFile(io.BytesIO(data)) as archive:
+            entries = archive.infolist()
+    except zipfile.BadZipFile as exc:
+        raise ValueError("Document is not a valid ZIP-based file") from exc
+
+    if len(entries) > MAX_ARCHIVE_ENTRIES:
+        raise ValueError("Document has too many parts to extract")
+    total = sum(entry.file_size for entry in entries)
+    if total > MAX_ARCHIVE_UNCOMPRESSED_BYTES:
+        raise ValueError("Document is too large to extract once uncompressed")
+    for entry in entries:
+        if entry.file_size > MAX_COMPRESSION_RATIO * max(entry.compress_size, 1):
+            raise ValueError(f"Document part {entry.filename!r} has a suspicious compression ratio")
+
+
+def _as_text(result: Any) -> str | None:
+    if result is None:
+        return None
+    if isinstance(result, list | tuple):
+        result = "\n\n".join(str(part) for part in result)
+    text = str(result)
+    if len(text) > MAX_EXTRACTED_CHARS:
+        logger.warning(
+            "Extracted text truncated from %d to %d chars", len(text), MAX_EXTRACTED_CHARS
+        )
+        text = text[:MAX_EXTRACTED_CHARS]
+    return text
+
+
+class DocumentExtractor:
+    """Wraps textxtract AsyncTextExtractor for API-side document extraction.
+
+    Examples::
+
+        extractor = DocumentExtractor()
+        text = await extractor.extract(pdf_bytes, "report.pdf")
+    """
+
+    def __init__(self, extractor: Any | None = None):
+        if extractor is not None:
+            self.extractor = extractor
+            return
+
+        if AsyncTextExtractor is None:
+            raise ImportError(
+                "textxtract is required for document extraction. "
+                'Install with `pip install "10xgraph-api[media]"`'
+            )
+
+        self.extractor = AsyncTextExtractor()
+
+    async def extract(
+        self,
+        data: bytes | str,
+        filename: str | None = None,
+        mime_type: str | None = None,
+    ) -> str | None:
+        """Extract text from bytes or a local path.
+
+        For bytes, the handler comes from ``mime_type`` when given, otherwise from the
+        extension of ``filename``, and only document types are extracted: archives and
+        unknown types return ``None``. The client's file name is never handed to textxtract.
+
+        Args:
+            data: Raw bytes or local path.
+            filename: Required when passing bytes (used only when ``mime_type`` is absent).
+            mime_type: The declared content type; preferred for picking the handler.
+
+        Returns:
+            Extracted text (capped at ``MAX_EXTRACTED_CHARS``), or ``None`` when the file
+            type is not extracted.
+
+        Raises:
+            ValueError: For a missing filename, an oversized or bomb-like document, a
+                timeout, or failed extraction.
+        """
+        if isinstance(data, bytes) and not filename:
+            raise ValueError("filename must be provided when extracting from bytes")
+
+        try:
+            if isinstance(data, bytes):
+                extension = _handler_extension(filename, mime_type)
+                if extension is None:
+                    logger.info(
+                        "Not extracting %s (%s): not a supported document", filename, mime_type
+                    )
+                    return None
+                if extension == ".docx":
+                    _check_zip_container(data)
+                result = await asyncio.wait_for(
+                    self.extractor.extract(data, f"upload{extension}"),
+                    timeout=EXTRACTION_TIMEOUT_SECONDS,
+                )
+                return _as_text(result)
+            return _as_text(await self.extractor.extract(data))
+
+        except TimeoutError as exc:
+            logger.warning("Document extraction timed out for %s", filename)
+            raise ValueError("Document extraction timed out") from exc
+
+        except FileTypeNotSupportedError:  # type: ignore
+            logger.warning("Document type not supported for extraction: %s", filename)
+            return None
+        except ExtractionError as exc:  # type: ignore
+            logger.exception("Document extraction failed for %s", filename)
+            raise ValueError("Failed to extract text from document") from exc