text2cypher 0.1.0__py3-none-any.whl

text2cypher/_results.py

@@ -0,0 +1,91 @@
+"""What came of one question: the result, its attempts and what a query returned."""
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+from text2cypher._validation import ValidationResult
+from text2cypher.llm import Usage
+
+type SchemaScope = Literal["filtered", "full"]
+type Status = Literal["valid", "unanswerable", "failed"]
+
+
+@dataclass(frozen=True, slots=True)
+class Attempt:
+    """One query the LLM wrote, how it was validated, and what writing it cost.
+
+    ``validation`` is ``None`` when the query is empty (the LLM declined), which is
+    not validated.
+
+    Attributes:
+        query: The query, empty if the LLM declined.
+        validation: How the query was validated, or ``None`` for an empty query.
+        schema_scope: The schema the query was written for: ``"filtered"`` or ``"full"``.
+        usage: The tokens the call that wrote the query used, if the client reported them.
+    """
+
+    query: str
+    validation: ValidationResult | None
+    schema_scope: SchemaScope
+    usage: Usage | None
+
+
+@dataclass(frozen=True, slots=True)
+class Records:
+    """What a query returned: the names of its columns, and its rows.
+
+    Values are the client's own: nodes, relationships and paths are ``falkordb``
+    objects.
+
+    Attributes:
+        columns: The names of the columns, as the query returns them.
+        rows: The rows, one tuple of values per row.
+    """
+
+    columns: tuple[str, ...]
+    rows: tuple[tuple[Any, ...], ...]
+
+
+@dataclass(frozen=True, slots=True)
+class Result:
+    """What came of one question: every attempt, in order, and how it ended.
+
+    ``valid``: the last query passed validation. ``unanswerable``: the LLM declined,
+    so there is no query. ``failed``: the corrections ran out; the last attempt has
+    the last query and its issues.
+
+    When execution is on and the status is ``valid``, ``records`` has what the query
+    returned, or ``execution_error`` says why FalkorDB did not run it (the status
+    stays ``valid``: the query passed validation).
+
+    A result is hashable unless it holds ``records``, whose values are the client's own
+    objects.
+
+    Attributes:
+        question: The question that was asked.
+        status: How it ended: ``"valid"``, ``"unanswerable"`` or ``"failed"``.
+        attempts: Every query the LLM wrote, in order, with its validation, schema and usage.
+        records: What the final query returned, when it was executed and ran.
+        execution_error: Why the final query did not run, when it was executed and failed.
+    """
+
+    question: str
+    status: Status
+    attempts: tuple[Attempt, ...]
+    records: Records | None = None
+    execution_error: str | None = None
+
+    @property
+    def query(self) -> str:
+        """The last query, empty if the LLM declined."""
+        return self.attempts[-1].query
+
+    @property
+    def usage(self) -> Usage | None:
+        """The tokens of all attempts, or ``None`` if any attempt did not report them."""
+        total = Usage(0, 0)
+        for attempt in self.attempts:
+            if attempt.usage is None:
+                return None
+            total += attempt.usage
+        return total

text2cypher/_schema_source.py

@@ -0,0 +1,60 @@
+"""Where the snapshot of a graph comes from: a new capture, or one that is kept."""
+
+import time
+from datetime import timedelta
+
+from cypher_guardrail import AsyncLiveSnapshot, Snapshot, SnapshotSource
+from falkordb.asyncio import FalkorDB
+
+from text2cypher._errors import GraphNotFoundError
+
+# A live snapshot refreshes itself in the background after this long; ours is refreshed by
+# ``SchemaSource.read`` when it is due, so it must never be due on its own.
+_NEVER = timedelta(days=36500)
+
+
+class SchemaSource:
+    """Reads the snapshot of a graph that the prompt and the validators are made from.
+
+    Args:
+        client: A ``falkordb.asyncio`` client.
+        refresh_seconds: ``None`` captures a new snapshot on every read. A number keeps one
+            per graph and, when that many seconds have passed since it was checked, checks
+            it again: a few cheap queries, and a new capture only if the graph changed.
+    """
+
+    def __init__(self, client: FalkorDB, refresh_seconds: float | None) -> None:
+        self._client = client
+        self._refresh_seconds = refresh_seconds
+        # graph name -> its live snapshot and when it was last checked (monotonic seconds)
+        self._live: dict[str, tuple[AsyncLiveSnapshot, float]] = {}
+
+    async def read(self, graph: str) -> tuple[Snapshot, SnapshotSource]:
+        """The graph's snapshot, and the source the validators should use.
+
+        Raises:
+            GraphNotFoundError: If the graph does not exist, or has no labelled nodes or
+                relationships. FalkorDB has no graph without data, so this is a missing
+                graph or a typo. Nothing is kept for it, so one created later is found.
+        """
+        snapshot, source = await self._snapshot(graph)
+        if not snapshot.labels and not snapshot.relationship_types:
+            self._live.pop(graph, None)
+            raise GraphNotFoundError(
+                f"The graph {graph!r} does not exist, or has no labelled nodes or relationships"
+            )
+        return snapshot, source
+
+    async def _snapshot(self, graph: str) -> tuple[Snapshot, SnapshotSource]:
+        if self._refresh_seconds is None:
+            snapshot = await Snapshot.capture_async(self._client, graph, types=True)
+            return snapshot, snapshot
+        live, checked = self._live.get(graph) or (
+            AsyncLiveSnapshot(self._client, graph, refresh_every=_NEVER, types=True),
+            -float("inf"),
+        )
+        if time.monotonic() - checked >= self._refresh_seconds:
+            await live.refresh()
+            checked = time.monotonic()
+        self._live[graph] = (live, checked)
+        return live.current(), live

text2cypher/_validation/__init__.py

@@ -0,0 +1,15 @@
+from text2cypher._validation.issues import (
+    Issue,
+    Source,
+    ValidationResult,
+    drop_repeated_property_issues,
+)
+from text2cypher._validation.validator import QueryValidator
+
+__all__ = [
+    "Issue",
+    "QueryValidator",
+    "Source",
+    "ValidationResult",
+    "drop_repeated_property_issues",
+]

text2cypher/_validation/issues.py

@@ -0,0 +1,100 @@
+"""The problems a validator found in a query, and how overlapping ones are merged."""
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any, Literal, cast
+
+from cypher_guardrail.issues import (
+    INVALID_PROPERTY_ACCESS,
+    UNKNOWN_LABEL,
+    UNKNOWN_REL_TYPE,
+)
+from cypher_guardrail.issues import Issue as GuardrailIssue
+
+type Source = Literal["syntax", "schema", "properties"]
+
+
+@dataclass(frozen=True, slots=True)
+class Issue:
+    """A problem found by one of the cypher-guardrail validators.
+
+    ``raw`` is cypher-guardrail's own typed issue: read its structured fields
+    (``label``, ``suggested_path``, ...) and never parse ``description``. It is a dict,
+    so it is left out of the hash (equal issues still have equal hashes).
+
+    Attributes:
+        source: The validator that found it: ``"syntax"``, ``"schema"`` or ``"properties"``.
+        raw: cypher-guardrail's issue: a ``code``, a ``description`` and fields that depend
+            on the code.
+    """
+
+    source: Source
+    raw: GuardrailIssue = field(hash=False)
+
+    @property
+    def fields(self) -> Mapping[str, Any]:
+        """``raw`` as a plain mapping, for reading the fields that depend on the code."""
+        return cast("Mapping[str, Any]", self.raw)
+
+    @property
+    def code(self) -> str:
+        """The issue's code, e.g. ``cypher_guardrail.Warning.UnknownLabel``."""
+        return self.raw["code"]
+
+    @property
+    def description(self) -> str:
+        """The issue in words, for people; read the fields of ``raw`` instead of parsing it."""
+        return self.raw["description"]
+
+
+@dataclass(frozen=True, slots=True)
+class ValidationResult:
+    """The outcome of validating one query against a graph.
+
+    A score is ``None`` when it carries no information: the properties score when
+    the query reads no property, and either score when the syntax is invalid and
+    that validator found nothing. Validity never depends on the scores.
+
+    Attributes:
+        issues: Every issue found, in the order syntax, schema, properties.
+        syntax_ok: Whether FalkorDB accepts the query.
+        schema_score: The schema validator's score, 0 to 1, or ``None``.
+        properties_score: The properties validator's score, 0 to 1, or ``None``.
+    """
+
+    issues: tuple[Issue, ...]
+    syntax_ok: bool
+    schema_score: float | None
+    properties_score: float | None
+
+    @property
+    def valid(self) -> bool:
+        """Whether the syntax is fine and no validator reported an issue."""
+        return self.syntax_ok and not self.issues
+
+
+def drop_repeated_property_issues(issues: list[Issue]) -> list[Issue]:
+    """Drop property issues about a label or relationship type already reported unknown.
+
+    cypher-guardrail reports a property of an unknown label, or of an unknown
+    relationship type, as missing too, so one root cause would otherwise appear twice.
+    A property issue says whether its name is a node label or a relationship type, so a
+    label and a relationship type that share a name are not mixed up.
+    """
+    unknown = {
+        ("node", issue.raw["label"])
+        for issue in issues
+        if issue.raw["code"] == UNKNOWN_LABEL
+    } | {
+        ("relationship", issue.raw["relationship_type"])
+        for issue in issues
+        if issue.raw["code"] == UNKNOWN_REL_TYPE
+    }
+    return [
+        issue
+        for issue in issues
+        if not (
+            issue.raw["code"] == INVALID_PROPERTY_ACCESS
+            and (issue.raw["kind"], issue.raw["name"]) in unknown
+        )
+    ]

text2cypher/_validation/validator.py

@@ -0,0 +1,106 @@
+"""Validate a Cypher query against a FalkorDB graph with cypher-guardrail."""
+
+import asyncio
+
+from cypher_guardrail import (
+    AsyncPropertiesValidator,
+    AsyncSchemaValidator,
+    AsyncSyntaxValidator,
+    SnapshotSource,
+)
+from cypher_guardrail import GraphNotFoundError as GuardrailGraphNotFoundError
+from cypher_guardrail.backend import get_async_backend
+from falkordb.asyncio import FalkorDB
+from falkordb.asyncio.graph import AsyncGraph
+
+from text2cypher._errors import GraphNotFoundError
+from text2cypher._validation.issues import (
+    Issue,
+    Source,
+    ValidationResult,
+    drop_repeated_property_issues,
+)
+
+
+class QueryValidator:
+    """Runs cypher-guardrail's syntax, schema and properties validators on a query.
+
+    All three always run, even when the syntax is invalid: they can still report
+    real problems elsewhere in the query. The query is valid when the syntax is ok
+    and no validator reported an issue; scores are never used to decide that. A
+    query that can write to the graph is never valid (cypher-guardrail's default).
+
+    Args:
+        client: A ``falkordb.asyncio`` client, or one of its graphs. The client is
+            bound to one event loop, so use one validator per loop.
+        max_concurrency: The most requests in flight at once.
+        snapshot: A snapshot (or a live one) of the graph. The validators then answer
+            from it and ask the database only to plan the query and to confirm a problem
+            before reporting it, so a query that is fine costs one request. The answers are
+            the same as without it, unless the graph changed since the snapshot.
+    """
+
+    def __init__(
+        self,
+        client: FalkorDB | AsyncGraph,
+        *,
+        max_concurrency: int = 8,
+        snapshot: SnapshotSource | None = None,
+    ) -> None:
+        # One backend shared by the validators: one snapshot cache, one concurrency limit
+        self._backend = get_async_backend(client, max_concurrency, snapshot=snapshot)
+        self._syntax = AsyncSyntaxValidator(self._backend)
+        self._schema = AsyncSchemaValidator(self._backend)
+        self._properties = AsyncPropertiesValidator(self._backend)
+
+    async def validate(self, query: str, graph: str | None = None) -> ValidationResult:
+        """Validate ``query`` against ``graph`` (optional when given a graph).
+
+        Raises:
+            ValueError: If there is no graph to validate against.
+            GraphNotFoundError: If the graph does not exist.
+        """
+        self._backend.check_database(graph)
+        try:
+            with self._backend.snapshot():
+                try:
+                    # Unlike gather, a failure cancels the other validators before the
+                    # snapshot is released
+                    async with asyncio.TaskGroup() as group:
+                        syntax = group.create_task(self._syntax.validate(query, graph))
+                        schema = group.create_task(self._schema.validate(query, graph))
+                        properties = group.create_task(
+                            self._properties.validate(query, graph)
+                        )
+                except ExceptionGroup as failures:
+                    # Callers should see the validator's own error, not a group
+                    raise failures.exceptions[0] from None
+        except GuardrailGraphNotFoundError as error:
+            # The graph can go away after its schema was read
+            raise GraphNotFoundError(str(error)) from error
+        syntax_ok, syntax_issues = syntax.result()
+        schema_score, schema_issues = schema.result()
+        properties_score, properties_issues = properties.result()
+
+        issues = [
+            *_wrap("syntax", syntax_issues),
+            *_wrap("schema", schema_issues),
+            *_wrap("properties", properties_issues),
+        ]
+        return ValidationResult(
+            issues=tuple(drop_repeated_property_issues(issues)),
+            syntax_ok=syntax_ok,
+            schema_score=_meaningful(schema_score, syntax_ok, schema_issues),
+            properties_score=_meaningful(
+                properties_score, syntax_ok, properties_issues
+            ),
+        )
+
+
+def _wrap(source: Source, raw_issues: list) -> list[Issue]:
+    return [Issue(source=source, raw=raw) for raw in raw_issues]
+
+
+def _meaningful(score: float | None, syntax_ok: bool, raw_issues: list) -> float | None:
+    """Drop a score from a validator that found nothing in an invalid query."""
+    return score if syntax_ok or raw_issues else None

text2cypher/_workflow.py

@@ -0,0 +1,99 @@
+"""The loop: generate a query, validate it, correct it, and escalate the schema."""
+
+from collections.abc import Awaitable, Callable
+from dataclasses import dataclass
+
+from text2cypher._config import Config
+from text2cypher._generation import Generation, correct_cypher, generate_cypher
+from text2cypher._results import Attempt, Result, SchemaScope
+from text2cypher._validation import Issue, ValidationResult
+from text2cypher.llm import AsyncLLMClient
+from text2cypher.schema import GraphSchema, SchemaFilter
+
+
+@dataclass(frozen=True, slots=True)
+class _Failed:
+    """A query that did not pass validation, and why, to be corrected."""
+
+    query: str
+    issues: tuple[Issue, ...]
+
+
+async def run_workflow(
+    question: str,
+    schema: GraphSchema,
+    *,
+    llm: AsyncLLMClient,
+    validate: Callable[[str], Awaitable[ValidationResult]],
+    config: Config,
+    schema_filter: SchemaFilter | None = None,
+) -> Result:
+    """Answer ``question`` with a query that passes ``validate``, if one can be found.
+
+    The first query and the first ``config.attempts_w_filtered`` corrections use the
+    filtered schema, the later corrections the full one, until ``config.total_attempts``
+    corrections have been made. Without a ``schema_filter``, or when it keeps the whole
+    schema, there is no filtered schema and every attempt uses the full one.
+
+    An empty query is the LLM declining. With the filtered schema that may be the
+    filter's doing, so the same request is made again with the full schema, once and
+    without costing a correction. Declining on the full schema ends the loop.
+    """
+    filtered = None if schema_filter is None else schema_filter(schema, question)
+    if filtered == schema:
+        filtered = None
+    attempts: list[Attempt] = []
+    # What the next request corrects; None for the first query
+    failed: _Failed | None = None
+    corrections = 0
+    # The LLM declined on the filtered schema, which is no longer used
+    escalated = False
+
+    while True:
+        scope, used = _choose_schema(
+            schema, filtered, config, corrections=corrections, escalated=escalated
+        )
+        generation = await _write(llm, question, used, failed)
+        if not generation.query:
+            attempts.append(Attempt("", None, scope, generation.usage))
+            if scope == "full":
+                return Result(question, "unanswerable", tuple(attempts))
+            escalated = True
+            continue
+
+        validation = await validate(generation.query)
+        attempts.append(Attempt(generation.query, validation, scope, generation.usage))
+        if validation.valid:
+            return Result(question, "valid", tuple(attempts))
+        if corrections == config.total_attempts:
+            return Result(question, "failed", tuple(attempts))
+        corrections += 1
+        failed = _Failed(generation.query, validation.issues)
+
+
+def _choose_schema(
+    schema: GraphSchema,
+    filtered: GraphSchema | None,
+    config: Config,
+    *,
+    corrections: int,
+    escalated: bool,
+) -> tuple[SchemaScope, GraphSchema]:
+    """The filtered schema while there is one, the budget for it lasts and the LLM has
+    not declined with it; the full schema otherwise."""
+    if (
+        filtered is not None
+        and not escalated
+        and corrections <= config.attempts_w_filtered
+    ):
+        return "filtered", filtered
+    return "full", schema
+
+
+async def _write(
+    llm: AsyncLLMClient, question: str, schema: GraphSchema, failed: _Failed | None
+) -> Generation:
+    """The first query, or a correction of the one that failed."""
+    if failed is None:
+        return await generate_cypher(llm, schema, question)
+    return await correct_cypher(llm, schema, question, failed.query, failed.issues)

text2cypher/llm/__init__.py

@@ -0,0 +1,5 @@
+from text2cypher.llm._base import AsyncLLMClient, Completion, Usage
+from text2cypher.llm._fake import Call, FakeLLM
+from text2cypher.llm._openai import OpenAIClient
+
+__all__ = ["AsyncLLMClient", "Call", "Completion", "FakeLLM", "OpenAIClient", "Usage"]

text2cypher/llm/_base.py

@@ -0,0 +1,43 @@
+"""What the rest of the library needs from an LLM."""
+
+from dataclasses import dataclass
+from typing import NamedTuple, Protocol
+
+
+@dataclass(frozen=True, slots=True)
+class Usage:
+    """The tokens one LLM call used; ``+`` adds two.
+
+    Attributes:
+        input_tokens: Tokens in the messages sent.
+        output_tokens: Tokens in the reply.
+    """
+
+    input_tokens: int
+    output_tokens: int
+
+    def __add__(self, other: "Usage") -> "Usage":
+        return Usage(
+            self.input_tokens + other.input_tokens,
+            self.output_tokens + other.output_tokens,
+        )
+
+
+class Completion(NamedTuple):
+    """One reply: its text, and the tokens it used when the provider says so.
+
+    Attributes:
+        text: The reply, with the query in a fenced block (an empty block declines).
+        usage: The tokens the call used, or ``None`` if they are not known.
+    """
+
+    text: str
+    usage: Usage | None
+
+
+class AsyncLLMClient(Protocol):
+    """A stateless chat model: every call is one system and one user message."""
+
+    async def complete(self, system: str, user: str) -> Completion:
+        """Ask the model, with a system message and a user message, and return its reply."""
+        ...

text2cypher/llm/_fake.py

@@ -0,0 +1,47 @@
+"""A scripted LLM for tests."""
+
+from collections.abc import Iterable
+from typing import NamedTuple
+
+from text2cypher.llm._base import Completion
+
+
+class Call(NamedTuple):
+    """What an LLM was asked once.
+
+    Attributes:
+        system: The system message.
+        user: The user message.
+    """
+
+    system: str
+    user: str
+
+
+class FakeLLM:
+    """Answers with the given replies in order, and records every call.
+
+    A reply is a ``Completion``, or a plain string for one without usage.
+
+    Attributes:
+        calls: What it was asked, in order.
+    """
+
+    def __init__(self, replies: Iterable[str | Completion]) -> None:
+        self._replies = iter(replies)
+        self.calls: list[Call] = []
+
+    async def complete(self, system: str, user: str) -> Completion:
+        """Record the call and return the next reply.
+
+        Raises:
+            RuntimeError: If every reply has been used.
+        """
+        self.calls.append(Call(system, user))
+        try:
+            reply = next(self._replies)
+        except StopIteration:
+            raise RuntimeError(
+                f"FakeLLM has no reply left for call {len(self.calls)}"
+            ) from None
+        return reply if isinstance(reply, Completion) else Completion(reply, None)

text2cypher/llm/_openai.py

@@ -0,0 +1,59 @@
+"""An LLM client for OpenAI's Chat Completions API (and compatible servers)."""
+
+from openai import AsyncOpenAI, omit
+from openai.types.chat import ChatCompletionMessageParam
+
+from text2cypher.llm._base import Completion, Usage
+
+
+class OpenAIClient:
+    """Asks an OpenAI chat model, one system and one user message per call.
+
+    A model that has no system role (some served by compatible servers) needs
+    ``system_role=False``: the system text then goes first in the one user message.
+
+    Args:
+        model: The model to ask. There is no default: the right one depends on the
+            account and on the schema being prompted.
+        client: The ``AsyncOpenAI`` to use, e.g. one with another ``base_url`` for a
+            compatible server. By default one is made from the environment.
+        temperature: ``None`` leaves it out, for models that only accept their own.
+        system_role: ``False`` sends one user message, the system text followed by the
+            user text, for a model that has no system role.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        *,
+        client: AsyncOpenAI | None = None,
+        temperature: float | None = 0,
+        system_role: bool = True,
+    ) -> None:
+        self._model = model
+        self._client = client or AsyncOpenAI()
+        self._temperature = omit if temperature is None else temperature
+        self._system_role = system_role
+
+    async def complete(self, system: str, user: str) -> Completion:
+        """Ask the model, and return its reply and, when the API reports it, the tokens used."""
+        response = await self._client.chat.completions.create(
+            model=self._model,
+            messages=self._messages(system, user),
+            temperature=self._temperature,
+        )
+        used = response.usage
+        return Completion(
+            text=response.choices[0].message.content or "",
+            usage=None
+            if used is None
+            else Usage(used.prompt_tokens, used.completion_tokens),
+        )
+
+    def _messages(self, system: str, user: str) -> list[ChatCompletionMessageParam]:
+        if self._system_role:
+            return [
+                {"role": "system", "content": system},
+                {"role": "user", "content": user},
+            ]
+        return [{"role": "user", "content": f"{system}\n\n{user}"}]

text2cypher/schema/__init__.py

@@ -0,0 +1,14 @@
+from text2cypher.schema._filter import SchemaFilter, filter_schema
+from text2cypher.schema._from_snapshot import schema_from_snapshot
+from text2cypher.schema._model import GraphSchema, Index, Path
+from text2cypher.schema._render import render_schema
+
+__all__ = [
+    "GraphSchema",
+    "Index",
+    "Path",
+    "SchemaFilter",
+    "filter_schema",
+    "render_schema",
+    "schema_from_snapshot",
+]