Graph Database (Apache AGE)¶

async def close(self) -> None:
    """Close the connection pool and release models."""
    await self.pool.close()
    embeddings.close()
    self.opened = False

async def create(
    self,
    node: GraphModelT,
) -> GraphModelT:
    """Create a node and its relationships in the graph."""
    await self._execute_batch(cypher.create(node))
    if _embeddable_descriptors(type(node)):
        await self._auto_embed(node)
    return node

async def delete(self, node: models.GraphModel) -> None:
    """Delete a node, its relationships, and embeddings.

    The Cypher delete runs via AGE (requires autocommit),
    then embeddings are cleaned up on the same connection.

    """
    self._require_open()
    stmt = cypher.delete(node)

    async def run(conn: psycopg.AsyncConnection[typing.Any]) -> None:
        await self._execute_on(
            conn,
            stmt.cypher,
            stmt.params,
        )
        if _embeddable_descriptors(type(node)):
            await self._delete_embeddings_where(
                conn,
                node_label=type(node).__name__,
                node_id=node.id,
            )

    await self._run_retrying_poisoned(run)

async def delete_node_embeddings(
    self,
    node_label: str,
    node_id: str,
) -> None:
    """Delete every embedding row for a node.

    The counterpart to ``embed_node`` for callers that delete
    nodes with raw Cypher instead of ``delete``.

    """
    self._require_open()
    async with self.pool.connection() as conn:
        await self._delete_embeddings_where(
            conn,
            node_label=node_label,
            node_id=node_id,
        )

async def embed_node(
    self,
    node: models.GraphModel,
    *,
    raise_on_error: bool = False,
) -> None:
    """(Re)generate the embeddings for *node*.

    ``create``/``merge``/``delete`` keep embeddings in sync on
    their own, but ``execute`` takes raw Cypher and has no model
    to inspect, so callers that write nodes that way must call
    this after the write.

    Failures are logged and swallowed like the automatic path,
    which suits a best-effort call alongside a graph write. Pass
    ``raise_on_error`` when the embedding *is* the job (a reindex)
    and a silent failure would be reported as success.

    """
    if raise_on_error:
        await self._embed_fields(node)
    else:
        await self._auto_embed(node)

async def execute(
    self,
    query_template: str,
    params: dict[str, typing.Any] | None = None,
    columns: list[str] | None = None,
    raw: bool = False,
) -> list[dict[str, typing.Any]]:
    """Wrap a Cypher query in SQL and execute it.

    Parameters in *params* are serialized via
    ``_cypher_param()`` using Cypher-compatible escaping
    and interpolated into *query_template* via
    ``sql.SQL.format()``.  Templated callers double their
    literal braces (``{{``/``}}``) and reference params as
    ``{name}``.

    Pass *raw* ``True`` to execute *query_template* as
    literal Cypher with no ``format()`` pass — required for
    arbitrary user-entered queries whose map literals use
    single braces.  *params* is ignored when *raw* is set.

    The Cypher query is wrapped in AGE's ``cypher()``
    function.  *columns* defines the ``AS (...)`` clause
    — pass one name per value in the Cypher ``RETURN``
    clause.  Defaults to ``['n']`` for single-column
    returns.

    """
    self._require_open()

    async def run(
        conn: psycopg.AsyncConnection[typing.Any],
    ) -> list[dict[str, typing.Any]]:
        return await self._execute_on(
            conn,
            query_template,
            params,
            columns,
            raw,
        )

    return await self._run_retrying_poisoned(run)

async def locked(
    self,
    key: str,
    work: collections.abc.Callable[
        [CypherExecute], collections.abc.Awaitable[RunResultT]
    ],
) -> RunResultT:
    """Run *work* in one transaction under an advisory lock on *key*.

    *work* gets an ``execute`` function with the arguments of
    :meth:`execute`. Its queries run in the transaction. The lock
    is transaction-scoped (``pg_advisory_xact_lock``), so calls with
    the same *key* run one at a time. Each statement after the lock
    sees the commits of the calls before it. An exception in *work*
    rolls the transaction back and releases the lock.

    """
    self._require_open()

    async def run(conn: psycopg.AsyncConnection[typing.Any]) -> RunResultT:
        async with conn.transaction():
            await conn.execute(
                'SELECT pg_advisory_xact_lock(hashtextextended(%s, 0))',
                (key,),
            )
            return await work(functools.partial(self._execute_on, conn))

    return await self._run_retrying_poisoned(run)

async def match(
    self,
    node_type: type[ModelT],
    params: dict[str, typing.Any] | None = None,
    order_by: str | None = None,
) -> list[ModelT]:
    """Match nodes and return model instances.

    Deserialization prefers ``model_validate`` (so field
    validators run) and falls back to ``model_construct``
    when validation fails.

    """
    stmt = cypher.match(node_type, params, order_by)
    raw_rows = await self.execute(
        stmt.cypher,
        stmt.params,
    )
    results: list[ModelT] = []
    for row in raw_rows:
        for value in row.values():
            props = parse_agtype(value)
            if isinstance(props, dict):
                results.append(
                    self._row_to_model(node_type, props),
                )
    return results

async def merge(
    self,
    node: GraphModelT,
    match_on: list[str] | None = None,
) -> GraphModelT:
    """Upsert a node and its relationships in the graph."""
    await self._execute_batch(
        cypher.merge(node, match_on),
    )
    if _embeddable_descriptors(type(node)):
        await self._auto_embed(node)
    return node

async def open(self) -> None:
    """Open the connection pool."""
    await self.pool.open()
    self.opened = True

async def search(
    self,
    query: str,
    *,
    model_name: str = 'text',
    node_label: str | None = None,
    attribute: str | None = None,
    node_ids: collections.abc.Collection[str] | None = None,
    limit: int = 10,
    distance_threshold: float | None = None,
) -> list[SearchResult]:
    """Search for nodes by semantic similarity.

    Embeds *query* using the specified model, then
    performs a cosine similarity search against the
    ``embeddings`` table.  Results are ordered by
    distance (ascending = most similar).

    When *node_ids* is provided the search is restricted to those
    node ids in SQL (``node_id = ANY(...)``), so callers can scope
    results (e.g. to a single organization) without fetching and
    post-filtering a wider result set. An empty collection matches
    nothing.

    """
    self._require_open()
    vector = await embeddings.aembed_one(
        query,
        model_name,
    )
    vec = sql.Placeholder('vec')
    dims = embeddings.get_dimensions(model_name)
    distance = sql.SQL(
        '(embedding::vector({dims})) <=> ({vec}::vector({dims}))',
    ).format(
        dims=sql.Literal(dims),
        vec=vec,
    )
    params: dict[str, typing.Any] = {
        'vec': vector,
        'model': model_name,
        'limit': limit,
    }
    query_sql = sql.SQL(
        'SELECT node_label, node_id, attribute,'
        '       chunk_text,'
        '       {distance} AS distance'
        '  FROM public.embeddings'
        ' WHERE model_name = {model}'
    ).format(
        distance=distance,
        model=sql.Placeholder('model'),
    )
    if node_label is not None:
        params['label'] = node_label
        query_sql += sql.SQL(
            ' AND node_label = {label}',
        ).format(label=sql.Placeholder('label'))
    if attribute is not None:
        params['attribute'] = attribute
        query_sql += sql.SQL(
            ' AND attribute = {attribute}',
        ).format(attribute=sql.Placeholder('attribute'))
    if node_ids is not None:
        params['node_ids'] = list(node_ids)
        query_sql += sql.SQL(
            ' AND node_id = ANY({node_ids})',
        ).format(node_ids=sql.Placeholder('node_ids'))
    if distance_threshold is not None:
        params['threshold'] = distance_threshold
        query_sql += sql.SQL(
            ' AND {distance} <= {threshold}',
        ).format(
            distance=distance,
            threshold=sql.Placeholder('threshold'),
        )
    query_sql += sql.SQL(
        ' ORDER BY {distance} LIMIT {limit}',
    ).format(
        distance=distance,
        limit=sql.Placeholder('limit'),
    )
    async with self.pool.connection() as conn:
        async with conn.cursor(
            row_factory=rows.dict_row,
        ) as cur:
            await cur.execute(query_sql, params)
            result_rows = await cur.fetchall()
    return [
        SearchResult(
            node_label=r['node_label'],
            node_id=r['node_id'],
            attribute=r['attribute'],
            chunk_text=r['chunk_text'],
            distance=r['distance'],
        )
        for r in result_rows
    ]

async def search_nodes(
    self,
    node_type: type[ModelT],
    query: str,
    *,
    model_name: str = 'text',
    limit: int = 10,
) -> list[ModelT]:
    """Search and return full node instances.

    Combines vector search with graph node retrieval.
    Results are deduplicated by ``id`` (multiple chunks
    from the same node may match).

    """
    # Over-fetch embedding rows so deduplication still
    # yields enough distinct nodes for the requested limit.
    chunk_multiplier = 5
    results = await self.search(
        query,
        model_name=model_name,
        node_label=node_type.__name__,
        limit=limit * chunk_multiplier,
    )
    node_ids = list(
        dict.fromkeys(r.node_id for r in results),
    )[:limit]
    if not node_ids:
        return []
    label = node_type.__name__
    id_list = ', '.join(f'{{{f"id{i}"}}}' for i in range(len(node_ids)))
    cypher_q = f'MATCH (n:{label}) WHERE n.id IN [{id_list}] RETURN n'
    params: dict[str, typing.Any] = {
        f'id{i}': nid for i, nid in enumerate(node_ids)
    }
    raw_rows = await self.execute(cypher_q, params)
    # Re-order to match the ranking from search()
    by_id: dict[str, ModelT] = {}
    for row in raw_rows:
        for value in row.values():
            props = parse_agtype(value)
            if isinstance(props, dict):
                nid = props.get('id')
                node = self._row_to_model(node_type, props)
                if nid is not None:
                    by_id[nid] = node
    return [by_id[nid] for nid in node_ids if nid in by_id]