lucid-graphql 0.2.0__py3-none-any.whl

lucid/__init__.py

@@ -0,0 +1,67 @@
+"""lucid — natural language to GraphQL, driven by an agentic workflow."""
+
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version as _version
+
+from lucid.client import (
+    EventStream as EventStream,
+)
+from lucid.client import (
+    LucidClient as LucidClient,
+)
+from lucid.client import (
+    QueryStream as QueryStream,
+)
+from lucid.client import (
+    ask as ask,
+)
+from lucid.client import (
+    ask_stream as ask_stream,
+)
+from lucid.client import (
+    create as create,
+)
+from lucid.client import (
+    generate as generate,
+)
+from lucid.client import (
+    generate_stream as generate_stream,
+)
+from lucid.errors import (
+    LucidError as LucidError,
+)
+from lucid.errors import (
+    QueryExecutionError as QueryExecutionError,
+)
+from lucid.errors import (
+    QueryValidationError as QueryValidationError,
+)
+from lucid.errors import (
+    SchemaError as SchemaError,
+)
+from lucid.errors import (
+    TransportError as TransportError,
+)
+from lucid.models import (
+    Event as Event,
+)
+from lucid.models import (
+    Response as Response,
+)
+from lucid.models import (
+    RunMetrics as RunMetrics,
+)
+from lucid.transport import (
+    ExecutionResult as ExecutionResult,
+)
+from lucid.transport import (
+    HttpxTransport as HttpxTransport,
+)
+from lucid.transport import (
+    Transport as Transport,
+)
+
+try:
+    __version__ = _version("lucid-graphql")
+except PackageNotFoundError:  # running from a source tree without an installed dist
+    __version__ = "0.0.0+unknown"

lucid/agent.py

@@ -0,0 +1,302 @@
+"""The Strands agent that authors a query, plus the guardrails around it.
+
+``stream_workflow`` is the single instrumented core: it yields progress
+``Event``s while the agent works and returns the final workflow state. The
+non-streaming entry points simply drain it, so the two can never diverge.
+"""
+
+import asyncio
+import contextvars
+import logging
+import queue
+import threading
+from collections import deque
+from collections.abc import AsyncGenerator, AsyncIterator, Generator, Iterator
+from pathlib import Path
+from typing import Any, cast
+
+from graphql import GraphQLSchema
+from strands import Agent
+from strands.hooks import AfterToolCallEvent, BeforeToolCallEvent, HookProvider, HookRegistry
+from strands.models import Model
+from strands.types.exceptions import EventLoopException
+
+from lucid.errors import LucidError, QueryExecutionError, QueryValidationError
+from lucid.models import Event, RunMetrics
+from lucid.tools import WorkflowState, build_tools
+from lucid.transport import Transport
+
+logger = logging.getLogger("lucid")
+
+_SYSTEM_PROMPT = """\
+You write a single valid GraphQL query that answers the user's question.
+
+Use the tools to explore the schema before writing the query. Never guess type
+or field names — look them up with list_root_fields, search_schema and get_type.
+
+Rules:
+- Start from list_root_fields to find the right entry point.
+- Select only the fields needed to answer the question; do not over-fetch.
+- Fields whose type is an interface or union require inline fragments
+  (`... on ConcreteType`) for type-specific fields.
+- You MUST call {goal_tool} with your query and only finish once it succeeds.
+  If it reports errors, fix the query and try again.
+- When done, reply with exactly the final query text and nothing else.
+"""
+
+_NUDGE = (
+    "You have not yet produced a query that {goal}. "
+    "Use the tools to fix the query, call {goal_tool} until it succeeds, "
+    "then reply with the final query text."
+)
+
+
+def _system_prompt(goal_tool: str, instructions: str | None) -> str:
+    prompt = _SYSTEM_PROMPT.format(goal_tool=goal_tool)
+    if instructions:
+        prompt += f"\nAdditional guidance from the application:\n{instructions}\n"
+    return prompt
+
+
+class _IterationGuard(HookProvider):
+    """Stops the run with a typed error once the correction cap is exhausted."""
+
+    def __init__(self, state: WorkflowState, max_iterations: int) -> None:
+        self._state = state
+        self._max_iterations = max_iterations
+
+    def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
+        registry.add_callback(AfterToolCallEvent, self._after_tool_call)
+
+    def _after_tool_call(self, event: AfterToolCallEvent) -> None:
+        if isinstance(event.exception, LucidError):
+            # The tool executor re-invokes this hook with our own raise attached;
+            # re-raising is what makes it propagate out of the event loop.
+            raise event.exception
+        if event.exception is not None:
+            return
+        if self._state.validation_failures >= self._max_iterations:
+            raise QueryValidationError(
+                f"no valid query after {self._state.validation_failures} validation attempts",
+                query=self._state.last_query,
+                errors=self._state.last_errors,
+            )
+        if self._state.execution_failures >= self._max_iterations:
+            raise QueryExecutionError(
+                f"execution still failing after {self._state.execution_failures} attempts",
+                query=self._state.last_query,
+                errors=self._state.last_errors,
+            )
+
+
+def _tool_message(name: str, tool_input: dict[str, Any]) -> str:
+    match name:
+        case "list_root_fields":
+            return "listing the root query fields"
+        case "search_schema":
+            return f"searching the schema for {tool_input.get('keyword', '?')!r}"
+        case "get_type":
+            return f"inspecting type {tool_input.get('name', '?')}"
+        case "validate_query":
+            return "validating a candidate query"
+        case "execute_query":
+            return "executing the query"
+    return f"calling {name}"
+
+
+class _EventCollector(HookProvider):
+    """Translates tool-call hooks into ``tool`` and ``attempt`` events."""
+
+    _ATTEMPT_TOOLS = ("validate_query", "execute_query")
+
+    def __init__(self, state: WorkflowState, pending: deque[Event]) -> None:
+        self._state = state
+        self._pending = pending
+        self._attempts = 0
+
+    def register_hooks(self, registry: HookRegistry, **kwargs: Any) -> None:
+        registry.add_callback(BeforeToolCallEvent, self._before_tool_call)
+        registry.add_callback(AfterToolCallEvent, self._after_tool_call)
+
+    @staticmethod
+    def _tool_input(event: BeforeToolCallEvent | AfterToolCallEvent) -> dict[str, Any]:
+        tool_input = event.tool_use.get("input")
+        return tool_input if isinstance(tool_input, dict) else {}
+
+    def _before_tool_call(self, event: BeforeToolCallEvent) -> None:
+        name = event.tool_use["name"]
+        self._pending.append(
+            Event(kind="tool", message=_tool_message(name, self._tool_input(event)))
+        )
+
+    def _after_tool_call(self, event: AfterToolCallEvent) -> None:
+        if event.exception is not None or event.tool_use["name"] not in self._ATTEMPT_TOOLS:
+            return
+        self._attempts += 1
+        query = self._tool_input(event).get("query")
+        if self._state.last_errors:
+            error = "; ".join(self._state.last_errors)
+            message = f"attempt {self._attempts} failed: {error}"
+        else:
+            error = None
+            verb = "executed" if event.tool_use["name"] == "execute_query" else "validated"
+            message = f"attempt {self._attempts}: query {verb} successfully"
+        self._pending.append(
+            Event(
+                kind="attempt",
+                message=message,
+                iteration=self._attempts,
+                query=query,
+                error=error,
+            )
+        )
+
+
+def _metrics(agent: Agent, state: WorkflowState) -> RunMetrics:
+    usage = agent.event_loop_metrics.accumulated_usage
+    return RunMetrics(
+        input_tokens=usage["inputTokens"],
+        output_tokens=usage["outputTokens"],
+        total_tokens=usage["totalTokens"],
+        cycles=agent.event_loop_metrics.cycle_count,
+        iterations=state.validation_failures + state.execution_failures,
+    )
+
+
+def _goal_reached(state: WorkflowState, execute: bool) -> bool:
+    if execute:
+        return state.executed_query is not None
+    return state.validated_query is not None
+
+
+def _drain_async(items: AsyncIterator[dict[str, Any]]) -> Iterator[dict[str, Any]]:
+    """Minimal sync bridge over the agent's async event stream (prototype).
+
+    The whole stream is consumed by a single task on a dedicated thread and
+    items are handed back over a queue. Step-driving the generator with one
+    ``run_until_complete`` per item would resume it in a fresh Task (and thus
+    a fresh contextvars Context) each time, which corrupts context-scoped
+    state such as OpenTelemetry spans. The caller's context is carried into
+    the thread so tracing stays parented correctly.
+    """
+    results: queue.Queue[tuple[str, Any]] = queue.Queue()
+    stop = threading.Event()
+    context = contextvars.copy_context()
+
+    async def consume() -> None:
+        try:
+            async for item in items:
+                results.put(("item", item))
+                if stop.is_set():
+                    break
+        finally:
+            if isinstance(items, AsyncGenerator):
+                await items.aclose()
+
+    def run() -> None:
+        try:
+            # asyncio.run (vs a bare event loop) matters: it cancels pending
+            # tasks and finalizes abandoned async generators (e.g. provider
+            # SDK streams) via shutdown_asyncgens before closing the loop.
+            asyncio.run(consume())
+            results.put(("done", None))
+        except BaseException as error:  # handed to the consuming thread to re-raise
+            results.put(("error", error))
+
+    thread = threading.Thread(target=lambda: context.run(run), name="lucid-stream", daemon=True)
+    thread.start()
+    try:
+        while True:
+            kind, value = results.get()
+            if kind == "done":
+                return
+            if kind == "error":
+                raise value
+            yield value
+    finally:
+        stop.set()
+        thread.join()
+
+
+def _stream_turn(agent: Agent, prompt: str, pending: deque[Event]) -> Generator[Event]:
+    """Run one agent turn, yielding thinking/tool/attempt events as they happen."""
+    thinking: list[str] = []
+
+    def flush_thinking() -> Iterator[Event]:
+        text = "".join(thinking).strip()
+        thinking.clear()
+        if text:
+            yield Event(kind="thinking", message=text)
+
+    def flush_pending() -> Iterator[Event]:
+        yield from flush_thinking()
+        while pending:
+            yield pending.popleft()
+
+    stream = _drain_async(aiter(agent.stream_async(prompt)))
+    try:
+        for item in stream:
+            if pending:
+                yield from flush_pending()
+            if "reasoningText" in item:
+                thinking.append(str(item["reasoningText"]))
+            elif "data" in item:
+                thinking.append(str(item["data"]))
+    except EventLoopException as error:
+        if isinstance(error.original_exception, LucidError):
+            raise error.original_exception from error
+        raise
+    yield from flush_pending()
+
+
+def stream_workflow(
+    question: str,
+    *,
+    schema: GraphQLSchema,
+    sdl_path: Path,
+    model: Model,
+    max_iterations: int,
+    transport: Transport | None = None,
+    instructions: str | None = None,
+) -> Generator[Event, None, tuple[WorkflowState, RunMetrics]]:
+    """Run the agentic workflow for one question, yielding progress events.
+
+    With a ``transport`` the goal is a successfully executed query (ask mode);
+    without one it is a successfully validated query (generate mode). Optional
+    ``instructions`` are appended to the built-in system prompt to steer
+    generation (the workflow contract itself is not overridable). Returns
+    the final state and metrics via the generator's return value.
+    """
+    execute = transport is not None
+    goal_tool = "execute_query" if execute else "validate_query"
+    state = WorkflowState()
+    pending: deque[Event] = deque()
+    agent = Agent(
+        model=model,
+        tools=cast(list[Any], build_tools(schema, sdl_path, state, transport=transport)),
+        system_prompt=_system_prompt(goal_tool, instructions),
+        hooks=[_IterationGuard(state, max_iterations), _EventCollector(state, pending)],
+        callback_handler=None,
+    )
+
+    yield from _stream_turn(agent, question, pending)
+    nudges = 0
+    while not _goal_reached(state, execute) and nudges < max_iterations:
+        nudges += 1
+        logger.debug("goal not reached; nudging agent (attempt %d)", nudges)
+        goal = "executes successfully" if execute else "validates"
+        yield from _stream_turn(agent, _NUDGE.format(goal=goal, goal_tool=goal_tool), pending)
+
+    if not _goal_reached(state, execute):
+        if execute and state.validated_query is not None:
+            raise QueryExecutionError(
+                "agent finished without successfully executing a query",
+                query=state.last_query,
+                errors=state.last_errors,
+            )
+        raise QueryValidationError(
+            "agent finished without a validated query",
+            query=state.last_query,
+            errors=state.last_errors,
+        )
+    return state, _metrics(agent, state)

lucid/cli.py

@@ -0,0 +1,58 @@
+"""Command-line interface: ask a GraphQL endpoint a question.
+
+Runnable via ``uvx --from "lucid-graphql[anthropic]" lucid ...`` or after
+``uv tool install "lucid-graphql[anthropic]"``.
+"""
+
+import json
+
+import click
+
+from lucid.client import create
+from lucid.errors import LucidError
+from lucid.models import Event
+
+
+@click.command()
+@click.argument("question")
+@click.option("--url", "-u", required=True, help="The GraphQL endpoint.")
+@click.option(
+    "--generate",
+    "-g",
+    "generate_only",
+    is_flag=True,
+    help="Print the generated query instead of executing it.",
+)
+@click.option(
+    "--verbose",
+    "-v",
+    is_flag=True,
+    help="Write the agent's progress (thinking, tools, attempts) to stderr.",
+)
+def main(question: str, url: str, generate_only: bool, verbose: bool) -> None:
+    """Turn a natural-language QUESTION into a GraphQL query and run it against --url.
+
+    Prints the query result as JSON (or, with --generate, the query itself)
+    to stdout. Uses the default Anthropic model; set ANTHROPIC_API_KEY.
+    """
+    # ask/generate are just their stream drained, so the stream form with
+    # on_event=None is identical to the plain call — no need to branch.
+    on_event = _echo_progress if verbose else None
+    try:
+        client = create(url)
+        if generate_only:
+            click.echo(client.generate_stream(question, on_event=on_event).result())
+        else:
+            response = client.ask_stream(question, on_event=on_event).result()
+            click.echo(json.dumps(response.data, indent=2))
+    except LucidError as error:
+        click.echo(f"Error: {error}", err=True)
+        raise SystemExit(1) from error
+
+
+def _echo_progress(event: Event) -> None:
+    click.echo(event.message, err=True)
+
+
+if __name__ == "__main__":
+    main()