omagent 0.1.0__py3-none-any.whl

omagent/__init__.py

@@ -0,0 +1,21 @@
+"""omagent — testable tooling layer for LLM-assisted OpenModelica workflows."""
+
+from .errors import (
+    Diagnostic, Kind, Severity,
+    classify, parse_error_string, parse_ompython_exception, parse_simulation_messages, summarize_for_llm,
+)
+from .session import Backend, OMSession, OpResult
+from .loop import AgentLoop, Attempt, LLM, LoopResult, Verifier, extract_code, extract_model_name
+from .results import (SimulationResult, all_of, expect_bounds, expect_final,
+                      expect_value_at, load_result)
+
+__version__ = "0.1.0"
+__all__ = [
+    "Diagnostic", "Kind", "Severity", "classify", "parse_error_string",
+    "parse_ompython_exception", "parse_simulation_messages", "summarize_for_llm",
+    "Backend", "OMSession", "OpResult",
+    "AgentLoop", "Attempt", "LLM", "LoopResult", "Verifier",
+    "extract_code", "extract_model_name",
+    "SimulationResult", "load_result", "expect_final", "expect_value_at",
+    "expect_bounds", "all_of",
+]

omagent/errors.py

@@ -0,0 +1,259 @@
+"""Parser for OpenModelica compiler (omc) diagnostics.
+
+Turns raw getErrorString()/simulate() output into structured records that an
+LLM agent (or a human) can act on. This module is deliberately free of any
+OMPython or network dependency so it can be unit-tested anywhere.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field, asdict
+from enum import Enum
+from typing import Optional
+
+
+class Severity(str, Enum):
+    ERROR = "error"
+    WARNING = "warning"
+    NOTIFICATION = "notification"
+
+
+class Kind(str, Enum):
+    """Coarse classification of what went wrong — used to steer fix strategies."""
+
+    SYNTAX = "syntax"                # parse errors, missing tokens
+    LOOKUP = "lookup"                # class/variable not found
+    TYPE = "type"                    # type mismatches, unit issues
+    BALANCE = "balance"              # equation/variable count mismatch, singular systems
+    CONNECT = "connect"              # connector/connection problems
+    INITIALIZATION = "initialization"
+    RUNTIME = "runtime"              # simulation-time failures (solver, events, asserts)
+    OTHER = "other"
+
+
+# [/path/File.mo:12:3-14:20:writable] Error: message...
+_LOCATED = re.compile(
+    r"^\[(?P<file>[^\]]*?):(?P<l1>\d+):(?P<c1>\d+)-(?P<l2>\d+):(?P<c2>\d+):[^\]]*\]\s*"
+    r"(?P<sev>Error|Warning|Notification):\s*(?P<msg>.*)$",
+    re.DOTALL,
+)
+
+# Error: message...  (no source location)
+_BARE = re.compile(
+    r"^(?P<sev>Error|Warning|Notification):\s*(?P<msg>.*)$", re.DOTALL
+)
+
+_SEV_MAP = {
+    "Error": Severity.ERROR,
+    "Warning": Severity.WARNING,
+    "Notification": Severity.NOTIFICATION,
+}
+
+# Ordered: first match wins.
+_KIND_RULES: list[tuple[Kind, re.Pattern]] = [
+    (Kind.SYNTAX, re.compile(
+        r"Parse error|Parser error|Missing token|syntax error|unexpected token|"
+        r"Expected token", re.IGNORECASE)),
+    (Kind.LOOKUP, re.compile(
+        r"not found in scope|Class \S+ not found|Variable \S+ not found|"
+        r"Base class \S+ not found|component .* not found", re.IGNORECASE)),
+    (Kind.TYPE, re.compile(
+        r"Type mismatch|expected subtype|incompatible types|"
+        r"unit .* not compatible|Illegal type", re.IGNORECASE)),
+    (Kind.CONNECT, re.compile(
+        r"connector|connect\(|not a valid connect|flow variable|stream variable",
+        re.IGNORECASE)),
+    (Kind.BALANCE, re.compile(
+        r"imbalanced|structurally singular|too (?:many|few) equations|"
+        r"under-?determined|over-?determined|"
+        r"equations? \(\d+\).*variables? \(\d+\)|"
+        r"\d+ equations? and \d+ variables?", re.IGNORECASE)),
+    (Kind.INITIALIZATION, re.compile(
+        r"initial(?:ization| equation| conditions?)|start value|fixed=", re.IGNORECASE)),
+    (Kind.RUNTIME, re.compile(
+        r"Simulation execution failed|solver|integrator|assert\b|division by zero|"
+        r"nonlinear system|chattering|stopped at time|LOG_", re.IGNORECASE)),
+]
+
+
+@dataclass
+class Diagnostic:
+    severity: Severity
+    message: str
+    kind: Kind = Kind.OTHER
+    file: Optional[str] = None
+    line_start: Optional[int] = None
+    col_start: Optional[int] = None
+    line_end: Optional[int] = None
+    col_end: Optional[int] = None
+
+    def to_dict(self) -> dict:
+        d = asdict(self)
+        d["severity"] = self.severity.value
+        d["kind"] = self.kind.value
+        return d
+
+    def brief(self) -> str:
+        loc = ""
+        if self.file and self.line_start is not None:
+            loc = f"{self.file}:{self.line_start}: "
+        return f"[{self.severity.value}/{self.kind.value}] {loc}{self.message}"
+
+
+def classify(message: str) -> Kind:
+    for kind, pat in _KIND_RULES:
+        if pat.search(message):
+            return kind
+    return Kind.OTHER
+
+
+def _split_records(raw: str) -> list[str]:
+    """Split a getErrorString() blob into individual diagnostic records.
+
+    omc concatenates records as lines; a record starts with '[' (located) or a
+    severity keyword, and continuation lines belong to the previous record.
+    """
+    records: list[str] = []
+    current: list[str] = []
+    start = re.compile(r"^(\[|Error:|Warning:|Notification:)")
+    for line in raw.splitlines():
+        if start.match(line.strip()) and current:
+            records.append("\n".join(current).strip())
+            current = [line]
+        elif start.match(line.strip()):
+            current = [line]
+        elif current:
+            current.append(line)
+        elif line.strip():
+            # Leading unstructured text (e.g. simulation stdout) — keep as record.
+            current = [line]
+    if current:
+        records.append("\n".join(current).strip())
+    return [r for r in records if r.strip().strip('"')]
+
+
+def parse_error_string(raw: str) -> list[Diagnostic]:
+    """Parse the output of omc's getErrorString() into Diagnostics."""
+    if not raw:
+        return []
+    raw = raw.strip().strip('"').strip()
+    if not raw:
+        return []
+
+    out: list[Diagnostic] = []
+    for rec in _split_records(raw):
+        m = _LOCATED.match(rec)
+        if m:
+            msg = m.group("msg").strip()
+            out.append(Diagnostic(
+                severity=_SEV_MAP[m.group("sev")],
+                message=msg,
+                kind=classify(msg),
+                file=m.group("file") or None,
+                line_start=int(m.group("l1")),
+                col_start=int(m.group("c1")),
+                line_end=int(m.group("l2")),
+                col_end=int(m.group("c2")),
+            ))
+            continue
+        m = _BARE.match(rec)
+        if m:
+            msg = m.group("msg").strip()
+            out.append(Diagnostic(
+                severity=_SEV_MAP[m.group("sev")],
+                message=msg,
+                kind=classify(msg),
+            ))
+            continue
+        # Unstructured content (runtime stdout etc.)
+        out.append(Diagnostic(
+            severity=Severity.ERROR if "fail" in rec.lower() else Severity.NOTIFICATION,
+            message=rec.strip(),
+            kind=classify(rec),
+        ))
+    return out
+
+
+def parse_simulation_messages(messages: str) -> list[Diagnostic]:
+    """Parse the 'messages' field of an omc simulate() result record.
+
+    Runtime logs look like:  LOG_STDOUT | error | msg   or plain text lines.
+    """
+    if not messages:
+        return []
+    out: list[Diagnostic] = []
+    log_line = re.compile(
+        r"^\s*(?P<stream>LOG_\w+)\s*\|\s*(?P<level>\w+)\s*\|\s*(?P<msg>.*)$")
+    for line in messages.splitlines():
+        if not line.strip():
+            continue
+        m = log_line.match(line)
+        if m:
+            level = m.group("level").lower()
+            sev = (Severity.ERROR if level in ("error", "assert")
+                   else Severity.WARNING if level == "warning"
+                   else Severity.NOTIFICATION)
+            msg = m.group("msg").strip()
+            out.append(Diagnostic(severity=sev, message=msg,
+                                  kind=classify(msg) if classify(msg) != Kind.OTHER
+                                  else Kind.RUNTIME))
+        else:
+            sev = (Severity.ERROR if re.search(r"fail|error", line, re.IGNORECASE)
+                   else Severity.NOTIFICATION)
+            out.append(Diagnostic(severity=sev, message=line.strip(),
+                                  kind=classify(line)))
+    return out
+
+
+def summarize_for_llm(diags: list[Diagnostic], limit: int = 20) -> str:
+    """Compact, deduplicated plain-text summary suitable for an LLM fix prompt."""
+    errors = [d for d in diags if d.severity == Severity.ERROR]
+    warnings = [d for d in diags if d.severity == Severity.WARNING]
+    seen: set[str] = set()
+    lines: list[str] = []
+    for d in errors + warnings:
+        b = d.brief()
+        if b not in seen:
+            seen.add(b)
+            lines.append(b)
+        if len(lines) >= limit:
+            lines.append(f"... ({len(errors) + len(warnings) - limit} more suppressed)")
+            break
+    return "\n".join(lines) if lines else "No errors or warnings."
+
+
+# Newer OMPython raises OMCSessionException whose message embeds omc's log as
+# "[OMC log for 'sendExpression(...)']: [kind:level:id] message"
+_OMPY_EXC = re.compile(
+    r"\[(?P<kind>\w+):(?P<level>error|warning|notification):(?P<id>-?\d+)\]\s*"
+    r"(?P<msg>.*)", re.DOTALL)
+
+_OMPY_KIND_MAP = {
+    "syntax": Kind.SYNTAX,
+    "grammar": Kind.SYNTAX,
+    "simulation": Kind.RUNTIME,
+}
+
+
+def parse_ompython_exception(text: str) -> list[Diagnostic]:
+    """Parse an OMPython OMCSessionException message into Diagnostics.
+
+    Falls back to a single generic error diagnostic for unstructured messages
+    (e.g. ZeroMQ connection failures), so backend exceptions never vanish.
+    """
+    m = _OMPY_EXC.search(text or "")
+    if not m:
+        return [Diagnostic(Severity.ERROR, (text or "unknown backend failure").strip(),
+                           Kind.OTHER)]
+    level = m.group("level")
+    sev = (Severity.ERROR if level == "error"
+           else Severity.WARNING if level == "warning"
+           else Severity.NOTIFICATION)
+    msg = m.group("msg").strip()
+    # message-based classification wins when specific; else fall back to
+    # omc's own kind tag (syntax/grammar/simulation/...)
+    kind = classify(msg)
+    if kind == Kind.OTHER:
+        kind = _OMPY_KIND_MAP.get(m.group("kind").lower(), Kind.OTHER)
+    return [Diagnostic(sev, msg, kind)]

omagent/llm.py

@@ -0,0 +1,96 @@
+"""LLM adapters implementing the ``omagent.loop.LLM`` protocol.
+
+Currently ships an Anthropic adapter (``ClaudeLLM``). Any other provider can
+be used by implementing the one-method ``propose`` protocol; nothing in the
+loop depends on this module.
+
+Install: ``pip install omagent[llm]`` and set ``ANTHROPIC_API_KEY``.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+SYSTEM_PROMPT = (
+    "You are an expert Modelica engineer producing models for OpenModelica.\n"
+    "Rules:\n"
+    "- Return ONE complete, self-contained Modelica model, and nothing else: "
+    "no prose before or after. A ```modelica fence is acceptable.\n"
+    "- The model must simulate in plain OpenModelica. Prefer the Modelica "
+    "Standard Library; never invent classes.\n"
+    "- Give every state a start value with fixed=true unless the task says "
+    "otherwise, so initialization is deterministic.\n"
+    "- Use SI units and add unit attributes where natural.\n"
+    "- When shown a failed attempt and its errors, return the corrected FULL "
+    "model, not a diff or fragment."
+)
+
+_FRESH_TEMPLATE = """Task: write a Modelica model for the following.
+
+{task}
+
+Return the complete model."""
+
+_FIX_TEMPLATE = """Task: {task}
+
+The model below failed. Fix it and return the complete corrected model.
+
+Failed model:
+```modelica
+{code}
+```
+
+Diagnostics:
+{errors}"""
+
+
+class ClaudeLLM:
+    """Anthropic-backed implementation of the LLM protocol.
+
+    Records every round in ``self.transcript`` — these (task, code, errors,
+    response) tuples are exactly the material worth keeping as benchmark
+    seeds and paper data.
+    """
+
+    def __init__(self, model: str = "claude-sonnet-4-6", max_tokens: int = 3000,
+                 client: Optional[Any] = None, _force_import_error: bool = False):
+        if client is None or _force_import_error:
+            try:
+                if _force_import_error:
+                    raise ImportError
+                import anthropic
+                client = anthropic.Anthropic()
+            except ImportError as exc:
+                raise ImportError(
+                    "ClaudeLLM requires the anthropic package: "
+                    "pip install omagent[llm]") from exc
+        self.client = client
+        self.model = model
+        self.max_tokens = max_tokens
+        self.transcript: list[dict] = []
+
+    def propose(self, task: str, previous_code: Optional[str],
+                error_summary: Optional[str]) -> str:
+        if previous_code is None:
+            user = _FRESH_TEMPLATE.format(task=task)
+        else:
+            user = _FIX_TEMPLATE.format(
+                task=task, code=previous_code,
+                errors=error_summary or "(no diagnostics captured)")
+
+        msg = self.client.messages.create(
+            model=self.model,
+            max_tokens=self.max_tokens,
+            system=SYSTEM_PROMPT,
+            messages=[{"role": "user", "content": user}],
+        )
+        text = "\n".join(
+            block.text for block in msg.content
+            if getattr(block, "type", "") == "text")
+        self.transcript.append({
+            "task": task,
+            "previous_code": previous_code,
+            "error_summary": error_summary,
+            "response": text,
+        })
+        return text

omagent/loop.py

@@ -0,0 +1,188 @@
+"""Agentic generate -> load -> check -> simulate -> verify loop.
+
+The LLM is behind a narrow protocol so the loop is unit-testable with a fake
+and backend-agnostic in production (Anthropic, OpenAI, local models — anything
+that can implement `propose`).
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Callable, Optional, Protocol
+
+import difflib
+
+from .errors import Diagnostic, Kind, summarize_for_llm
+from .session import OMSession, OpResult
+
+
+class LLM(Protocol):
+    def propose(
+        self,
+        task: str,
+        previous_code: Optional[str],
+        error_summary: Optional[str],
+    ) -> str:
+        """Return complete Modelica code for `task`.
+
+        First call: previous_code/error_summary are None (fresh generation).
+        Later calls: both are set (repair a failed attempt).
+        """
+        ...
+
+
+# A verifier inspects the simulate() OpResult; returns None if satisfied,
+# else a human/LLM-readable complaint that is fed into the next fix round.
+Verifier = Callable[[OpResult], Optional[str]]
+
+_MODEL_NAME = re.compile(r"^\s*(?:model|block|package)\s+([A-Za-z_][A-Za-z0-9_]*)",
+                         re.MULTILINE)
+
+_CODE_FENCE = re.compile(r"```(?:modelica|mo)?\s*\n(.*?)```", re.DOTALL)
+
+
+def extract_code(llm_output: str) -> str:
+    """Accept raw Modelica or a fenced markdown block; return bare code."""
+    m = _CODE_FENCE.search(llm_output)
+    return (m.group(1) if m else llm_output).strip()
+
+
+def extract_model_name(code: str) -> Optional[str]:
+    m = _MODEL_NAME.search(code)
+    return m.group(1) if m else None
+
+
+@dataclass
+class Attempt:
+    n: int
+    code: str
+    stage: str                       # "load" | "check" | "simulate" | "verify" | "ok"
+    diagnostics: list[Diagnostic] = field(default_factory=list)
+    complaint: Optional[str] = None  # verifier feedback, if any
+
+    @property
+    def failed(self) -> bool:
+        return self.stage != "ok"
+
+
+@dataclass
+class LoopResult:
+    success: bool
+    attempts: list[Attempt]
+    model_name: Optional[str] = None
+    sim_result: Optional[OpResult] = None
+
+    @property
+    def final_code(self) -> Optional[str]:
+        return self.attempts[-1].code if self.attempts else None
+
+
+class AgentLoop:
+    def __init__(
+        self,
+        session: OMSession,
+        llm: LLM,
+        max_attempts: int = 4,
+        simulate_options: Optional[dict] = None,
+        verifier: Optional[Verifier] = None,
+    ):
+        if max_attempts < 1:
+            raise ValueError("max_attempts must be >= 1")
+        self.session = session
+        self.llm = llm
+        self.max_attempts = max_attempts
+        self.simulate_options = simulate_options or {}
+        self.verifier = verifier
+
+    def run(self, task: str, model_name: Optional[str] = None) -> LoopResult:
+        attempts: list[Attempt] = []
+        prev_code: Optional[str] = None
+        feedback: Optional[str] = None
+
+        for n in range(1, self.max_attempts + 1):
+            code = extract_code(self.llm.propose(task, prev_code, feedback))
+            name = model_name or extract_model_name(code)
+            if not name:
+                att = Attempt(n, code, "load")
+                att.complaint = ("Could not determine the model name: the code "
+                                 "must contain a top-level `model <Name>`.")
+                attempts.append(att)
+                prev_code, feedback = code, att.complaint
+                continue
+
+            att, sim = self._try_once(n, code, name)
+            attempts.append(att)
+            if not att.failed:
+                return LoopResult(True, attempts, name, sim)
+            prev_code = code
+            feedback = self._feedback(att)
+
+        return LoopResult(False, attempts, model_name or
+                          extract_model_name(attempts[-1].code))
+
+    # -- internals --------------------------------------------------------
+    def _try_once(self, n: int, code: str, name: str):
+        res = self.session.load_string(code)
+        if not res.success:
+            return Attempt(n, code, "load", res.diagnostics), None
+
+        res = self.session.check_model(name)
+        if not res.success:
+            return Attempt(n, code, "check", res.diagnostics), None
+
+        sim = self.session.simulate(name, **self.simulate_options)
+        if not sim.success:
+            return Attempt(n, code, "simulate", sim.diagnostics), None
+
+        if self.verifier is not None:
+            complaint = self.verifier(sim)
+            if complaint:
+                return Attempt(n, code, "verify", sim.diagnostics, complaint), sim
+
+        return Attempt(n, code, "ok", sim.diagnostics), sim
+
+    def _feedback(self, att: Attempt) -> str:
+        parts = [f"Attempt failed at stage '{att.stage}'."]
+        if att.diagnostics:
+            parts.append(summarize_for_llm(att.diagnostics))
+        parts.extend(self._lookup_suggestions(att.diagnostics))
+        if att.complaint:
+            parts.append(f"Verification feedback: {att.complaint}")
+        return "\n".join(parts)
+
+    _MISSING_CLASS = re.compile(r"(?:Class|Import)\s+([A-Za-z_][\w.]*\.[\w]+)"
+                                r"\s+not found")
+
+    def _lookup_suggestions(self, diags: list[Diagnostic],
+                            max_packages: int = 3) -> list[str]:
+        """For lookup failures, ask omc what the parent package really
+        contains and surface close matches — turns 'X not found' into
+        'did you mean RotationalEMF' (e.g. MSL 3.2 -> 4.x renames)."""
+        out: list[str] = []
+        seen: set[str] = set()
+        for d in diags:
+            if d.kind != Kind.LOOKUP:
+                continue
+            m = self._MISSING_CLASS.search(d.message)
+            if not m:
+                continue
+            fqn = m.group(1)
+            parent, _, leaf = fqn.rpartition(".")
+            if not parent or parent in seen:
+                continue
+            seen.add(parent)
+            names = self.session.class_names(parent)
+            if not names:
+                continue
+            close = difflib.get_close_matches(leaf, names, n=5, cutoff=0.4)
+            sub = [n for n in names
+                   if leaf.lower() in n.lower() and n not in close]
+            picks = (close + sub)[:6] or sorted(names)[:12]
+            out.append(
+                f"Hint: {fqn} does not exist in the loaded libraries. "
+                f"{parent} actually contains: {', '.join(picks)}. "
+                f"Use one of these exact names.")
+            if len(seen) >= max_packages:
+                break
+        return out