sim-atlas 0.0.1__py3-none-any.whl

sim_atlas/agent/__init__.py

@@ -0,0 +1,3 @@
+from ._runner import run_agent_stream
+
+__all__ = ["run_agent_stream"]

sim_atlas/agent/_prompt.py

@@ -0,0 +1,66 @@
+from ..models import AgentRequest, Annotation, GraphNodeContext
+from ..storage_interface import StorageInterface
+from .tools import get_tool_prompt_guidance_lines
+
+
+def _port_str(a: Annotation) -> str:
+    s = f"{a.label}:{a.datatype or '?'}"
+    if a.unit:
+        s += f"[{a.unit}]"
+    if a.quantity:
+        s += f"({a.quantity})"
+    if a.description:
+        s += f" — {a.description}"
+    return s
+
+
+def _node_to_context(node: GraphNodeContext) -> str:
+    inputs = ", ".join(_port_str(a) for a in node.inputs)
+    outputs = ", ".join(_port_str(a) for a in node.outputs)
+    desc = f" - {node.short_description}" if node.short_description else ""
+    atlas = f" (atlas_id={node.atlas_node_id})" if node.node_kind == "function" else ""
+    return f"  [{node.node_kind}:{node.graph_id}] {node.name}{atlas}{desc}\n    in: {inputs}\n    out: {outputs}"
+
+
+def build_system_prompt(request: AgentRequest, storage: StorageInterface) -> str:
+    node_lines = "\n".join(_node_to_context(n) for n in request.nodes)
+    edge_lines = "\n".join(
+        f"  {e.source_graph_id}/{e.source_handle} -> {e.target_graph_id}/{e.target_handle}"
+        for e in request.edges
+    )
+    opts = storage.get_filter_options()
+    filter_section = (
+        "## Available filter values\n"
+        "Use only values from these lists when supplying filter arguments to "
+        "find_compatible_nodes.\n"
+        f"datatypes: {', '.join(sorted(opts.datatypes)) or '(none)'}\n"
+        f"units: {', '.join(sorted(opts.units)) or '(none)'}\n"
+        f"quantities: {', '.join(sorted(opts.quantities)) or '(none)'}"
+    )
+    tool_guidance = "\n".join(f"- {line}" for line in get_tool_prompt_guidance_lines())
+
+    return f"""You are an expert workflow builder for scientific simulation pipelines.
+You have access to a catalog of simulation nodes (functions) that can be searched and connected.
+
+## Current graph
+
+### Nodes
+{node_lines or "  (empty)"}
+
+### Edges
+{edge_lines or "  (none)"}
+
+## Your role
+You are an editor of an existing workflow graph.
+Treat the current graph as intentional work that should be preserved unless change is necessary.
+Make the minimum changes needed to satisfy the user's request.
+Prefer adding new nodes and edges over modifying or removing existing ones.
+Only remove or rewire existing nodes when the user explicitly asks, or when it is genuinely necessary to produce a correct result.
+
+## Tools
+- If you face doubts, hard to solve problems, inconsistencies, or multiple options, ask the user a clarification question in plain text instead of guessing.
+{tool_guidance}
+- When you are finished, respond with a concise summary of only the changes you made.
+
+{filter_section}
+"""

sim_atlas/agent/_runner.py

@@ -0,0 +1,213 @@
+import json
+import logging
+import uuid
+from collections.abc import AsyncGenerator
+from typing import Any, cast
+
+from langfuse import Langfuse, propagate_attributes
+from langfuse.openai import AsyncOpenAI  # pyright: ignore[reportPrivateImportUsage]
+from openai.types.chat import (
+    ChatCompletionMessageParam,
+    ChatCompletionMessageToolCall,
+)
+from pydantic import ValidationError
+
+from ..models import AgentRequest
+from ..settings import load_settings
+from ..storage_interface import StorageInterface
+from ._prompt import build_system_prompt
+from ._sse import (
+    ErrorEvent,
+    GraphUpdateEvent,
+    MessageEvent,
+    ReasoningEvent,
+    ToolCallEvent,
+    ToolResultEvent,
+    TruncatedEvent,
+    ValidationEvent,
+    to_sse,
+)
+from .tools import TOOLS, ScratchGraph, ToolError, execute_tool, validate_graph
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.DEBUG)
+
+
+def _graph_update_event(scratch: ScratchGraph) -> GraphUpdateEvent:
+    return GraphUpdateEvent(
+        nodes=[n.model_dump(exclude_none=True) for n in scratch.nodes.values()],
+        edges=[e.model_dump() for e in scratch.edges],
+    )
+
+
+async def run_agent_stream(
+    request: AgentRequest,
+    storage: StorageInterface,
+) -> AsyncGenerator[str, None]:
+    """Async generator that streams SSE events while running the agent loop."""
+    settings = load_settings()
+    assert settings.llm_api_key and settings.llm_base_url and settings.llm_chat_model
+    lf = Langfuse(
+        public_key=settings.langfuse_public_key,
+        secret_key=settings.langfuse_secret_key,
+        host=settings.langfuse_host,
+        environment=settings.langfuse_environment,
+    )
+    client = AsyncOpenAI(api_key=settings.llm_api_key, base_url=settings.llm_base_url)
+    scratch = ScratchGraph(request.nodes, request.edges)
+
+    history_messages: list[ChatCompletionMessageParam] = [
+        cast(ChatCompletionMessageParam, {"role": m.role, "content": m.content})
+        for m in request.history
+    ]
+    messages: list[ChatCompletionMessageParam] = [
+        {"role": "system", "content": build_system_prompt(request, storage)},
+        *history_messages,
+        {"role": "user", "content": request.query},
+    ]
+
+    session_id = request.session_id or str(uuid.uuid4())
+    final_message: str | None = None
+    correction_rounds = 0
+    total_input_tokens = 0
+    total_output_tokens = 0
+    max_turns = settings.agent_max_iterations
+    with (
+        lf.start_as_current_observation(
+            as_type="agent",
+            name="workflow-builder",
+            input=request.query,
+        ) as root_span,
+        propagate_attributes(user_id=request.user_id, session_id=session_id),
+    ):
+        try:
+            # Runaway-check loop: exits naturally (break) when the agent finishes,
+            # or falls through (no break) when the turn limit is reached.
+            for _ in range(max_turns):
+                response = await client.chat.completions.create(
+                    model=settings.llm_chat_model,
+                    messages=messages,
+                    tools=TOOLS,
+                    tool_choice="auto",
+                )
+                choice = response.choices[0]
+                logger.debug(
+                    "LLM response message: %s",
+                    json.dumps(choice.message.model_dump(exclude_unset=True), indent=2),
+                )
+                if response.usage:
+                    total_input_tokens += response.usage.prompt_tokens
+                    total_output_tokens += response.usage.completion_tokens
+                messages.append(
+                    cast(
+                        ChatCompletionMessageParam,
+                        choice.message.model_dump(exclude_unset=True),
+                    )
+                )
+
+                if not choice.message.tool_calls:
+                    final_message = choice.message.content or "(no response)"
+                    with lf.start_as_current_observation(
+                        name="graph_validation"
+                    ) as span:
+                        validation_errors = validate_graph(scratch, storage)
+                        span.update(
+                            output={"errors": validation_errors},
+                        )
+                        if not validation_errors:
+                            break
+                    # Emit a validation event so the UI can show a correction round.
+                    yield to_sse(ValidationEvent(errors=validation_errors))
+                    correction_rounds += 1
+                    error_text = "\n".join(f"- {e}" for e in validation_errors)
+                    logger.debug(
+                        "Graph validation errors (stream); asking agent to correct:\n%s",
+                        error_text,
+                    )
+                    correction_message = (
+                        "The current graph has validation errors. "
+                        "Please fix them using the available tools:\n" + error_text
+                    )
+                    messages.append({"role": "user", "content": correction_message})
+                    continue
+
+                reasoning = (
+                    getattr(choice.message, "reasoning", None)
+                    or choice.message.content
+                    or None
+                )
+                if reasoning:
+                    yield to_sse(ReasoningEvent(content=reasoning))
+                for tc in choice.message.tool_calls:
+                    if not isinstance(tc, ChatCompletionMessageToolCall):
+                        continue
+                    args: dict[str, Any] = json.loads(tc.function.arguments)
+                    yield to_sse(ToolCallEvent(name=tc.function.name, args=args))
+
+                    try:
+                        result = await execute_tool(
+                            tc.function.name, args, storage, scratch
+                        )
+                        content = result
+                    except (ValidationError, ToolError) as exc:
+                        msg = (
+                            f"Invalid arguments for '{tc.function.name}': {exc}"
+                            if isinstance(exc, ValidationError)
+                            else str(exc)
+                        )
+                        content = json.dumps({"error": msg})
+
+                    yield to_sse(
+                        ToolResultEvent(name=tc.function.name, content=content)
+                    )
+                    messages.append(
+                        {
+                            "role": "tool",
+                            "tool_call_id": tc.id,
+                            "content": content,
+                        }
+                    )
+
+                yield to_sse(_graph_update_event(scratch))
+
+            truncated = final_message is None
+            if truncated:
+                # Turn limit reached without natural completion.
+                # Make one tool-free summary call so the history entry is honest.
+                messages.append(
+                    {
+                        "role": "user",
+                        "content": (
+                            "You have reached the maximum number of turns. "
+                            "Summarise in 2-3 sentences: what you have built so far "
+                            "and what still needs to be done to fulfil the original request. "
+                            "End with a clear statement of the next step the user should ask you to take."
+                        ),
+                    }
+                )
+
+                summary_response = await client.chat.completions.create(
+                    model=settings.llm_chat_model,
+                    messages=messages,
+                    tool_choice="none",
+                )
+                summary_choice = summary_response.choices[0]
+                final_message = summary_choice.message.content or "(turn limit reached)"
+                if summary_response.usage:
+                    total_input_tokens += summary_response.usage.prompt_tokens
+                    total_output_tokens += summary_response.usage.completion_tokens
+                logger.debug("Turn limit reached; summary: %s", final_message)
+
+            root_span.update(output=final_message)
+            yield to_sse(_graph_update_event(scratch))
+            yield to_sse(MessageEvent(content=final_message))
+            if truncated:
+                yield to_sse(TruncatedEvent())
+
+            lf.flush()
+        except Exception as exc:  # noqa: BLE001
+            logger.exception("Agent stream error")
+            root_span.update(level="ERROR", status_message=str(exc))
+            yield to_sse(ErrorEvent(message=str(exc)))
+        finally:
+            lf.flush()

sim_atlas/agent/_sse.py

@@ -0,0 +1,71 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel
+
+
+class ReasoningEvent(BaseModel):
+    type: Literal["reasoning"] = "reasoning"
+    content: str
+
+
+class ToolCallEvent(BaseModel):
+    type: Literal["tool_call"] = "tool_call"
+    name: str
+    args: dict[str, Any]
+
+
+class ToolResultEvent(BaseModel):
+    type: Literal["tool_result"] = "tool_result"
+    name: str
+    content: str
+
+
+class MessageEvent(BaseModel):
+    type: Literal["message"] = "message"
+    content: str
+
+
+class ClarificationEvent(BaseModel):
+    type: Literal["clarification"] = "clarification"
+    question: str
+    options: list[str]
+
+
+class ValidationEvent(BaseModel):
+    type: Literal["validation"] = "validation"
+    errors: list[str]
+
+
+class ErrorEvent(BaseModel):
+    type: Literal["error"] = "error"
+    message: str
+
+
+class GraphUpdateEvent(BaseModel):
+    type: Literal["graph_update"] = "graph_update"
+    nodes: list[dict[str, Any]]
+    edges: list[dict[str, Any]]
+
+
+class TruncatedEvent(BaseModel):
+    """Emitted when the runaway-check loop exhausts all turns without natural completion."""
+
+    type: Literal["truncated"] = "truncated"
+
+
+Event = (
+    ReasoningEvent
+    | ToolCallEvent
+    | ToolResultEvent
+    | MessageEvent
+    | ClarificationEvent
+    | ValidationEvent
+    | ErrorEvent
+    | GraphUpdateEvent
+    | TruncatedEvent
+)
+
+
+def to_sse(event: Event) -> str:
+    """Format an event payload as an SSE data frame."""
+    return f"data: {event.model_dump_json()}\n\n"

sim_atlas/agent/tools/__init__.py

@@ -0,0 +1,229 @@
+from dataclasses import dataclass
+from typing import Any
+
+from langfuse import get_client
+from openai.types.chat import ChatCompletionToolParam
+from pydantic import BaseModel
+
+from ...storage_interface import StorageInterface
+from ._errors import ToolError
+from ._graph import (
+    ADD_EDGE_DESCRIPTION,
+    ADD_FUNCTION_NODE_DESCRIPTION,
+    ADD_INPUT_NODE_DESCRIPTION,
+    ADD_OUTPUT_NODE_DESCRIPTION,
+    REMOVE_EDGE_DESCRIPTION,
+    REMOVE_NODE_DESCRIPTION,
+    AddEdgeInput,
+    AddEdgeResult,
+    AddFunctionNodeInput,
+    AddFunctionNodeResult,
+    AddInputNodeInput,
+    AddInputNodeResult,
+    AddOutputNodeInput,
+    AddOutputNodeResult,
+    RemoveEdgeInput,
+    RemoveEdgeResult,
+    RemoveNodeInput,
+    RemoveNodeResult,
+    ScratchGraph,
+    execute_add_edge,
+    execute_add_function_node,
+    execute_add_input_node,
+    execute_add_output_node,
+    execute_remove_edge,
+    execute_remove_node,
+    validate_graph,
+)
+from ._search import (
+    FIND_COMPATIBLE_DESCRIPTION,
+    GET_NODE_DETAILS_DESCRIPTION,
+    SEARCH_NODES_DESCRIPTION,
+    FindCompatibleNodesInput,
+    FindCompatibleNodesResult,
+    GetNodeDetailsInput,
+    GetNodeDetailsResult,
+    SearchNodesInput,
+    SearchNodesResult,
+    execute_find_compatible_nodes,
+    execute_get_node_details,
+    execute_search_nodes,
+)
+
+ExecutorFn = Any
+
+
+@dataclass(frozen=True)
+class ToolDef:
+    name: str
+    description: str
+    prompt_guidance: str
+    input_model: type[BaseModel]
+    output_model: type[BaseModel]
+    executor: ExecutorFn
+
+
+TOOL_REGISTRY: list[ToolDef] = [
+    ToolDef(
+        name="search_nodes",
+        description=SEARCH_NODES_DESCRIPTION,
+        prompt_guidance="Use search_nodes for intent-based discovery.",
+        input_model=SearchNodesInput,
+        output_model=SearchNodesResult,
+        executor=execute_search_nodes,
+    ),
+    ToolDef(
+        name="find_compatible_nodes",
+        description=FIND_COMPATIBLE_DESCRIPTION,
+        prompt_guidance=(
+            "Use find_compatible_nodes when you know a specific port signature "
+            "and want to find what connects to it."
+        ),
+        input_model=FindCompatibleNodesInput,
+        output_model=FindCompatibleNodesResult,
+        executor=execute_find_compatible_nodes,
+    ),
+    ToolDef(
+        name="get_node_details",
+        description=GET_NODE_DETAILS_DESCRIPTION,
+        prompt_guidance=(
+            "Use get_node_details when you need more information about a specific "
+            "node before deciding."
+        ),
+        input_model=GetNodeDetailsInput,
+        output_model=GetNodeDetailsResult,
+        executor=execute_get_node_details,
+    ),
+    ToolDef(
+        name="add_function_node",
+        description=ADD_FUNCTION_NODE_DESCRIPTION,
+        prompt_guidance=(
+            "Use add_function_node to add a node from the catalog; it returns "
+            "the assigned graph_id and port metadata — use that graph_id in "
+            "subsequent add_edge calls."
+        ),
+        input_model=AddFunctionNodeInput,
+        output_model=AddFunctionNodeResult,
+        executor=execute_add_function_node,
+    ),
+    ToolDef(
+        name="add_input_node",
+        description=ADD_INPUT_NODE_DESCRIPTION,
+        prompt_guidance=(
+            "Use add_input_node to expose a parameter or data source. "
+            "The single output port is always named 'output'."
+        ),
+        input_model=AddInputNodeInput,
+        output_model=AddInputNodeResult,
+        executor=execute_add_input_node,
+    ),
+    ToolDef(
+        name="add_output_node",
+        description=ADD_OUTPUT_NODE_DESCRIPTION,
+        prompt_guidance=(
+            "Use add_output_node to expose a result. "
+            "The single input port is always named 'input'."
+        ),
+        input_model=AddOutputNodeInput,
+        output_model=AddOutputNodeResult,
+        executor=execute_add_output_node,
+    ),
+    ToolDef(
+        name="add_edge",
+        description=ADD_EDGE_DESCRIPTION,
+        prompt_guidance=(
+            "Use add_edge to connect nodes; verify port names using the port "
+            "metadata returned by add_function_node or get_node_details."
+        ),
+        input_model=AddEdgeInput,
+        output_model=AddEdgeResult,
+        executor=execute_add_edge,
+    ),
+    ToolDef(
+        name="remove_edge",
+        description=REMOVE_EDGE_DESCRIPTION,
+        prompt_guidance=(
+            "Use remove_edge to disconnect two nodes by specifying all four "
+            "endpoint identifiers."
+        ),
+        input_model=RemoveEdgeInput,
+        output_model=RemoveEdgeResult,
+        executor=execute_remove_edge,
+    ),
+    ToolDef(
+        name="remove_node",
+        description=REMOVE_NODE_DESCRIPTION,
+        prompt_guidance=(
+            "Use remove_node to delete an existing node and all its connected edges."
+        ),
+        input_model=RemoveNodeInput,
+        output_model=RemoveNodeResult,
+        executor=execute_remove_node,
+    ),
+]
+
+_TOOL_MAP: dict[str, ToolDef] = {tool.name: tool for tool in TOOL_REGISTRY}
+
+
+def _schema_from_input_model(model: type[BaseModel]) -> dict[str, Any]:
+    schema = model.model_json_schema()
+    return {
+        "type": "object",
+        "properties": schema.get("properties", {}),
+        "required": schema.get("required", []),
+    }
+
+
+def _build_tools_list() -> list[ChatCompletionToolParam]:
+    return [
+        {
+            "type": "function",
+            "function": {
+                "name": tool.name,
+                "description": tool.description,
+                "parameters": _schema_from_input_model(tool.input_model),
+            },
+        }
+        for tool in TOOL_REGISTRY
+    ]
+
+
+TOOLS: list[ChatCompletionToolParam] = _build_tools_list()
+
+
+def get_tool_prompt_guidance_lines() -> list[str]:
+    return [tool.prompt_guidance for tool in TOOL_REGISTRY]
+
+
+async def execute_tool(
+    tool_name: str,
+    tool_args: dict[str, Any],
+    storage: StorageInterface,
+    scratch: ScratchGraph,
+) -> str:
+    lf = get_client()
+    with lf.start_as_current_observation(
+        as_type="tool",
+        name=tool_name,
+        input=tool_args,
+    ) as span:
+        tool = _TOOL_MAP.get(tool_name)
+        if tool is None:
+            raise ToolError(f"Unknown tool '{tool_name}'.")
+
+        validated_args = tool.input_model.model_validate(tool_args)
+        result = await tool.executor(validated_args, storage, scratch)
+        span.update(output=result)
+        return result
+
+
+__all__ = [
+    "TOOL_REGISTRY",
+    "TOOLS",
+    "ScratchGraph",
+    "ToolDef",
+    "ToolError",
+    "execute_tool",
+    "get_tool_prompt_guidance_lines",
+    "validate_graph",
+]

sim_atlas/agent/tools/_errors.py

@@ -0,0 +1,2 @@
+class ToolError(Exception):
+    """Raised when a tool execution fails in a recoverable way."""