cocoindex 0.3.4__cp311-abi3-manylinux_2_28_x86_64.whl

cocoindex/query_handler.py

@@ -0,0 +1,57 @@
+import dataclasses
+import numpy as np
+from numpy import typing as npt
+from typing import Generic, Any
+from .index import VectorSimilarityMetric
+import sys
+
+if sys.version_info >= (3, 13):
+    from typing import TypeVar
+else:
+    from typing_extensions import TypeVar  # PEP 696 backport
+
+
+@dataclasses.dataclass
+class QueryHandlerResultFields:
+    """
+    Specify field names in query results returned by the query handler.
+    This provides metadata for tools like CocoInsight to recognize structure of the query results.
+    """
+
+    embedding: list[str] = dataclasses.field(default_factory=list)
+    score: str | None = None
+
+
+@dataclasses.dataclass
+class QueryHandlerInfo:
+    """
+    Info to configure a query handler.
+    """
+
+    result_fields: QueryHandlerResultFields | None = None
+
+
+@dataclasses.dataclass
+class QueryInfo:
+    """
+    Info about the query.
+    """
+
+    embedding: list[float] | npt.NDArray[np.float32] | None = None
+    similarity_metric: VectorSimilarityMetric | None = None
+
+
+R = TypeVar("R", default=Any)
+
+
+@dataclasses.dataclass
+class QueryOutput(Generic[R]):
+    """
+    Output of a query handler.
+
+    results: list of results. Each result can be a dict or a dataclass.
+    query_info: Info about the query.
+    """
+
+    results: list[R]
+    query_info: QueryInfo = dataclasses.field(default_factory=QueryInfo)

cocoindex/runtime.py

@@ -0,0 +1,78 @@
+"""
+This module provides a standalone execution runtime for executing coroutines in a thread-safe
+manner.
+"""
+
+import threading
+import asyncio
+import inspect
+import warnings
+
+from typing import Any, Callable, Awaitable, TypeVar, Coroutine
+
+T = TypeVar("T")
+
+
+class _ExecutionContext:
+    _lock: threading.Lock
+    _event_loop: asyncio.AbstractEventLoop | None = None
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+
+    @property
+    def event_loop(self) -> asyncio.AbstractEventLoop:
+        """Get the event loop for the cocoindex library."""
+        with self._lock:
+            if self._event_loop is None:
+                loop = asyncio.new_event_loop()
+                self._event_loop = loop
+
+                def _runner(l: asyncio.AbstractEventLoop) -> None:
+                    asyncio.set_event_loop(l)
+                    l.run_forever()
+
+                threading.Thread(target=_runner, args=(loop,), daemon=True).start()
+            return self._event_loop
+
+    def run(self, coro: Coroutine[Any, Any, T]) -> T:
+        """Run a coroutine in the event loop, blocking until it finishes. Return its result."""
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+
+        loop = self.event_loop
+
+        if running_loop is not None:
+            if running_loop is loop:
+                raise RuntimeError(
+                    "CocoIndex sync API was called from inside CocoIndex's async context. "
+                    "Use the async variant of this method instead."
+                )
+            warnings.warn(
+                "CocoIndex sync API was called inside an existing event loop. "
+                "This may block other tasks. Prefer the async method.",
+                RuntimeWarning,
+                stacklevel=2,
+            )
+
+        fut = asyncio.run_coroutine_threadsafe(coro, loop)
+        try:
+            return fut.result()
+        except KeyboardInterrupt:
+            fut.cancel()
+            raise
+
+
+execution_context = _ExecutionContext()
+
+
+def to_async_call(call: Callable[..., Any]) -> Callable[..., Awaitable[Any]]:
+    if isinstance(call, (staticmethod, classmethod)):
+        base_call = call.__func__
+    else:
+        base_call = call
+    if inspect.iscoroutinefunction(base_call):
+        return call
+    return lambda *args, **kwargs: asyncio.to_thread(lambda: call(*args, **kwargs))

cocoindex/setting.py

@@ -0,0 +1,171 @@
+"""
+Data types for settings of the cocoindex library.
+"""
+
+import os
+
+from typing import Callable, Self, Any, overload
+from dataclasses import dataclass
+from . import _engine  # type: ignore
+
+
+def get_app_namespace(*, trailing_delimiter: str | None = None) -> str:
+    """Get the application namespace. Append the `trailing_delimiter` if not empty."""
+    app_namespace: str = _engine.get_app_namespace()
+    if app_namespace == "" or trailing_delimiter is None:
+        return app_namespace
+    return f"{app_namespace}{trailing_delimiter}"
+
+
+def split_app_namespace(full_name: str, delimiter: str) -> tuple[str, str]:
+    """Split the full name into the application namespace and the rest."""
+    parts = full_name.split(delimiter, 1)
+    if len(parts) == 1:
+        return "", parts[0]
+    return (parts[0], parts[1])
+
+
+@dataclass
+class DatabaseConnectionSpec:
+    """
+    Connection spec for relational database.
+    Used by both internal and target storage.
+    """
+
+    url: str
+    user: str | None = None
+    password: str | None = None
+    max_connections: int = 25
+    min_connections: int = 5
+
+
+@dataclass
+class GlobalExecutionOptions:
+    """Global execution options."""
+
+    # The maximum number of concurrent inflight requests, shared among all sources from all flows.
+    source_max_inflight_rows: int | None = 1024
+    source_max_inflight_bytes: int | None = None
+
+
+def _load_field(
+    target: dict[str, Any],
+    name: str,
+    env_name: str,
+    required: bool = False,
+    parse: Callable[[str], Any] | None = None,
+) -> None:
+    value = os.getenv(env_name)
+    if value is None:
+        if required:
+            raise ValueError(f"{env_name} is not set")
+    else:
+        if parse is None:
+            target[name] = value
+        else:
+            try:
+                target[name] = parse(value)
+            except Exception as e:
+                raise ValueError(
+                    f"failed to parse environment variable {env_name}: {value}"
+                ) from e
+
+
+@dataclass
+class Settings:
+    """Settings for the cocoindex library."""
+
+    database: DatabaseConnectionSpec | None = None
+    app_namespace: str = ""
+    global_execution_options: GlobalExecutionOptions | None = None
+
+    @classmethod
+    def from_env(cls) -> Self:
+        """Load settings from environment variables."""
+
+        database_url = os.getenv("COCOINDEX_DATABASE_URL")
+        if database_url is not None:
+            db_kwargs: dict[str, Any] = {"url": database_url}
+            _load_field(db_kwargs, "user", "COCOINDEX_DATABASE_USER")
+            _load_field(db_kwargs, "password", "COCOINDEX_DATABASE_PASSWORD")
+            _load_field(
+                db_kwargs,
+                "max_connections",
+                "COCOINDEX_DATABASE_MAX_CONNECTIONS",
+                parse=int,
+            )
+            _load_field(
+                db_kwargs,
+                "min_connections",
+                "COCOINDEX_DATABASE_MIN_CONNECTIONS",
+                parse=int,
+            )
+            database = DatabaseConnectionSpec(**db_kwargs)
+        else:
+            database = None
+
+        exec_kwargs: dict[str, Any] = dict()
+        _load_field(
+            exec_kwargs,
+            "source_max_inflight_rows",
+            "COCOINDEX_SOURCE_MAX_INFLIGHT_ROWS",
+            parse=int,
+        )
+        _load_field(
+            exec_kwargs,
+            "source_max_inflight_bytes",
+            "COCOINDEX_SOURCE_MAX_INFLIGHT_BYTES",
+            parse=int,
+        )
+        global_execution_options = GlobalExecutionOptions(**exec_kwargs)
+
+        app_namespace = os.getenv("COCOINDEX_APP_NAMESPACE", "")
+
+        return cls(
+            database=database,
+            app_namespace=app_namespace,
+            global_execution_options=global_execution_options,
+        )
+
+
+@dataclass
+class ServerSettings:
+    """Settings for the cocoindex server."""
+
+    # The address to bind the server to.
+    address: str = "127.0.0.1:49344"
+
+    # The origins of the clients (e.g. CocoInsight UI) to allow CORS from.
+    cors_origins: list[str] | None = None
+
+    @classmethod
+    def from_env(cls) -> Self:
+        """Load settings from environment variables."""
+        kwargs: dict[str, Any] = dict()
+        _load_field(kwargs, "address", "COCOINDEX_SERVER_ADDRESS")
+        _load_field(
+            kwargs,
+            "cors_origins",
+            "COCOINDEX_SERVER_CORS_ORIGINS",
+            parse=ServerSettings.parse_cors_origins,
+        )
+        return cls(**kwargs)
+
+    @overload
+    @staticmethod
+    def parse_cors_origins(s: str) -> list[str]: ...
+
+    @overload
+    @staticmethod
+    def parse_cors_origins(s: str | None) -> list[str] | None: ...
+
+    @staticmethod
+    def parse_cors_origins(s: str | None) -> list[str] | None:
+        """
+        Parse the CORS origins from a string.
+        """
+        return (
+            [o for e in s.split(",") if (o := e.strip()) != ""]
+            if s is not None
+            else None
+        )

cocoindex/setup.py

@@ -0,0 +1,92 @@
+"""
+This module provides APIs to manage the setup of flows.
+"""
+
+from . import setting
+from . import _engine  # type: ignore
+from .runtime import execution_context
+
+
+class SetupChangeBundle:
+    """
+    This class represents a bundle of setup changes.
+    """
+
+    _engine_bundle: _engine.SetupChangeBundle
+
+    def __init__(self, _engine_bundle: _engine.SetupChangeBundle):
+        self._engine_bundle = _engine_bundle
+
+    def __str__(self) -> str:
+        desc, _ = execution_context.run(self._engine_bundle.describe_async())
+        return desc  # type: ignore
+
+    def __repr__(self) -> str:
+        return self.__str__()
+
+    def apply(self, report_to_stdout: bool = False) -> None:
+        """
+        Apply the setup changes.
+        """
+        execution_context.run(self.apply_async(report_to_stdout=report_to_stdout))
+
+    async def apply_async(self, report_to_stdout: bool = False) -> None:
+        """
+        Apply the setup changes. Async version of `apply`.
+        """
+        await self._engine_bundle.apply_async(report_to_stdout=report_to_stdout)
+
+    def describe(self) -> tuple[str, bool]:
+        """
+        Describe the setup changes.
+        """
+        return execution_context.run(self.describe_async())  # type: ignore
+
+    async def describe_async(self) -> tuple[str, bool]:
+        """
+        Describe the setup changes. Async version of `describe`.
+        """
+        return await self._engine_bundle.describe_async()  # type: ignore
+
+    def describe_and_apply(self, report_to_stdout: bool = False) -> None:
+        """
+        Describe the setup changes and apply them if `report_to_stdout` is True.
+        Silently apply setup changes otherwise.
+        """
+        execution_context.run(
+            self.describe_and_apply_async(report_to_stdout=report_to_stdout)
+        )
+
+    async def describe_and_apply_async(self, *, report_to_stdout: bool = False) -> None:
+        """
+        Describe the setup changes and apply them if `report_to_stdout` is True.
+        Silently apply setup changes otherwise. Async version of `describe_and_apply`.
+        """
+        if report_to_stdout:
+            desc, is_up_to_date = await self.describe_async()
+            print("Setup status:\n")
+            print(desc)
+            if is_up_to_date:
+                print("No setup changes to apply.")
+                return
+        await self.apply_async(report_to_stdout=report_to_stdout)
+
+
+def flow_names_with_setup() -> list[str]:
+    """
+    Get the names of all flows that have been setup.
+    """
+    return execution_context.run(flow_names_with_setup_async())  # type: ignore
+
+
+async def flow_names_with_setup_async() -> list[str]:
+    """
+    Get the names of all flows that have been setup. Async version of `flow_names_with_setup`.
+    """
+    result = []
+    all_flow_names = await _engine.flow_names_with_setup_async()
+    for name in all_flow_names:
+        app_namespace, name = setting.split_app_namespace(name, ".")
+        if app_namespace == setting.get_app_namespace():
+            result.append(name)
+    return result

cocoindex/sources/__init__.py

@@ -0,0 +1,5 @@
+"""
+Sources supported by CocoIndex.
+"""
+
+from ._engine_builtin_specs import *

cocoindex/sources/_engine_builtin_specs.py

@@ -0,0 +1,120 @@
+"""All builtin sources."""
+
+from .. import op
+from ..auth_registry import TransientAuthEntryReference
+from ..setting import DatabaseConnectionSpec
+from dataclasses import dataclass
+import datetime
+
+
+class LocalFile(op.SourceSpec):
+    """Import data from local file system."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    path: str
+    binary: bool = False
+
+    # If provided, only files matching these patterns will be included.
+    # See https://docs.rs/globset/latest/globset/index.html#syntax for the syntax of the patterns.
+    included_patterns: list[str] | None = None
+
+    # If provided, files matching these patterns will be excluded.
+    # See https://docs.rs/globset/latest/globset/index.html#syntax for the syntax of the patterns.
+    excluded_patterns: list[str] | None = None
+
+    # If provided, files exceeding this size in bytes will be treated as non-existent.
+    max_file_size: int | None = None
+
+
+class GoogleDrive(op.SourceSpec):
+    """Import data from Google Drive."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    service_account_credential_path: str
+    root_folder_ids: list[str]
+    binary: bool = False
+    recent_changes_poll_interval: datetime.timedelta | None = None
+
+
+@dataclass
+class RedisNotification:
+    """Redis pub/sub configuration for event notifications."""
+
+    # Redis server URL (e.g., "redis://localhost:6379")
+    redis_url: str
+    # Redis channel name for pub/sub notifications
+    redis_channel: str
+
+
+class AmazonS3(op.SourceSpec):
+    """Import data from an Amazon S3 bucket. Supports optional prefix and file filtering by glob patterns."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    bucket_name: str
+    prefix: str | None = None
+    binary: bool = False
+    included_patterns: list[str] | None = None
+    excluded_patterns: list[str] | None = None
+    max_file_size: int | None = None
+    sqs_queue_url: str | None = None
+    redis: RedisNotification | None = None
+
+
+class AzureBlob(op.SourceSpec):
+    """
+    Import data from an Azure Blob Storage container. Supports optional prefix and file filtering by glob patterns.
+
+    Authentication mechanisms taken in the following order:
+    - SAS token (if provided)
+    - Account access key (if provided)
+    - Default Azure credential
+    """
+
+    _op_category = op.OpCategory.SOURCE
+
+    account_name: str
+    container_name: str
+    prefix: str | None = None
+    binary: bool = False
+    included_patterns: list[str] | None = None
+    excluded_patterns: list[str] | None = None
+
+    sas_token: TransientAuthEntryReference[str] | None = None
+    account_access_key: TransientAuthEntryReference[str] | None = None
+
+
+@dataclass
+class PostgresNotification:
+    """Notification for a PostgreSQL table."""
+
+    # Optional: name of the PostgreSQL channel to use.
+    # If not provided, will generate a default channel name.
+    channel_name: str | None = None
+
+
+class Postgres(op.SourceSpec):
+    """Import data from a PostgreSQL table."""
+
+    _op_category = op.OpCategory.SOURCE
+
+    # Table name to read from (required)
+    table_name: str
+
+    # Database connection reference (optional - uses default if not provided)
+    database: TransientAuthEntryReference[DatabaseConnectionSpec] | None = None
+
+    # Optional: specific columns to include (if None, includes all columns)
+    included_columns: list[str] | None = None
+
+    # Optional: column name to use for ordinal tracking (for incremental updates)
+    # Should be a timestamp, serial, or other incrementing column
+    ordinal_column: str | None = None
+
+    # Optional: when set, supports change capture from PostgreSQL notification.
+    notification: PostgresNotification | None = None
+
+    # Optional: SQL expression filter for rows (arbitrary SQL boolean expression)
+    filter: str | None = None