agents

__getattr__(name)

Dynamically import attributes on first access.

This avoids importing all agent modules at package import time, so a failure in one agent does not prevent using others.

Source code in src/ursa/agents/__init__.py

def __getattr__(name: str) -> Any:
    """Dynamically import attributes on first access.

    This avoids importing all agent modules at package import time,
    so a failure in one agent does not prevent using others.
    """
    try:
        module_name, attr_name = _lazy_attrs[name]
    except KeyError:
        raise AttributeError(
            f"module {__name__!r} has no attribute {name!r}"
        ) from None

    module = importlib.import_module(module_name, __name__)
    value = getattr(module, attr_name)
    # Cache the loaded attribute so subsequent access is fast
    globals()[name] = value
    return value

acquisition_agents

ArxivAgent

Bases: BaseAcquisitionAgent

Drop-in replacement for your existing ArxivAgent that reuses the generic flow. Keeps the same behaviors (download PDFs, image processing, summarization/RAG).

Source code in src/ursa/agents/acquisition_agents.py

class ArxivAgent(BaseAcquisitionAgent):
    """
    Drop-in replacement for your existing ArxivAgent that reuses the generic flow.
    Keeps the same behaviors (download PDFs, image processing, summarization/RAG).
    """

    def __init__(
        self,
        llm: BaseChatModel,
        *,
        process_images: bool = True,
        max_results: int = 3,
        download: bool = True,
        rag_embedding=None,
        database_path="arxiv_papers",
        summaries_path="arxiv_generated_summaries",
        vectorstore_path="arxiv_vectorstores",
        **kwargs,
    ):
        super().__init__(
            llm,
            rag_embedding=rag_embedding,
            process_images=process_images,
            max_results=max_results,
            database_path=database_path,
            summaries_path=summaries_path,
            vectorstore_path=vectorstore_path,
            download=download,
            **kwargs,
        )

    def _id(self, hit_or_item: dict[str, Any]) -> str:
        # hits from arXiv feed have 'id' like ".../abs/XXXX.YYYY"
        arxiv_id = hit_or_item.get("arxiv_id")
        if arxiv_id:
            return arxiv_id
        feed_id = hit_or_item.get("id", "")
        if "/abs/" in feed_id:
            return feed_id.split("/abs/")[-1]
        return _hash(json.dumps(hit_or_item))

    def _citation(self, item: ItemMetadata) -> str:
        return f"ArXiv ID: {item.get('id', '?')}"

    def _search(self, query: str) -> list[dict[str, Any]]:
        enc = quote(query)
        url = f"http://export.arxiv.org/api/query?search_query=all:{enc}&start=0&max_results={self.max_results}"
        try:
            resp = requests.get(url, timeout=15)
            resp.raise_for_status()
            feed = feedparser.parse(resp.content)
            entries = feed.entries if hasattr(feed, "entries") else []
            hits = []
            for e in entries:
                full_id = e.id.split("/abs/")[-1]
                hits.append({
                    "id": e.id,
                    "title": e.title.strip(),
                    "arxiv_id": full_id.split("/")[-1],
                })
            return hits
        except Exception as e:
            return [
                {
                    "id": _hash(query + str(time.time())),
                    "title": "Search error",
                    "error": str(e),
                }
            ]

    def _materialize(self, hit: dict[str, Any]) -> ItemMetadata:
        arxiv_id = self._id(hit)
        title = hit.get("title", "")
        pdf_url = f"https://arxiv.org/pdf/{arxiv_id}.pdf"
        local_path = os.path.join(self.database_path, f"{arxiv_id}.pdf")
        full_text = ""
        try:
            _download(pdf_url, local_path)
            full_text = read_pdf(local_path)
        except Exception as e:
            full_text = f"[Error loading ArXiv {arxiv_id}: {e}]"
        full_text = self._postprocess_text(full_text, local_path)
        return {
            "id": arxiv_id,
            "title": title,
            "url": pdf_url,
            "local_path": local_path,
            "full_text": full_text,
        }

BaseAcquisitionAgent

Bases: BaseAgent

A generic "acquire-then-summarize-or-RAG" agent.

Subclasses must implement

• _search(self, query) -> List[dict-like]: lightweight hits
• _materialize(self, hit) -> ItemMetadata: download or scrape and return populated item
• _id(self, hit_or_item) -> str: stable id for caching/file naming
• _citation(self, item) -> str: human-readable citation string

Optional hooks

• _postprocess_text(self, text, local_path) -> str (e.g., image interpretation)
• _filter_hit(self, hit) -> bool

Source code in src/ursa/agents/acquisition_agents.py

class BaseAcquisitionAgent(BaseAgent):
    """
    A generic "acquire-then-summarize-or-RAG" agent.

    Subclasses must implement:
      - _search(self, query) -> List[dict-like]: lightweight hits
      - _materialize(self, hit) -> ItemMetadata: download or scrape and return populated item
      - _id(self, hit_or_item) -> str: stable id for caching/file naming
      - _citation(self, item) -> str: human-readable citation string

    Optional hooks:
      - _postprocess_text(self, text, local_path) -> str (e.g., image interpretation)
      - _filter_hit(self, hit) -> bool
    """

    def __init__(
        self,
        llm: BaseChatModel,
        *,
        summarize: bool = True,
        rag_embedding=None,
        process_images: bool = True,
        max_results: int = 5,
        database_path: str = "acq_db",
        summaries_path: str = "acq_summaries",
        vectorstore_path: str = "acq_vectorstores",
        num_threads: int = 4,
        download: bool = True,
        **kwargs,
    ):
        super().__init__(llm, **kwargs)
        self.summarize = summarize
        self.rag_embedding = rag_embedding
        self.process_images = process_images
        self.max_results = max_results
        self.database_path = self.den / database_path
        self.summaries_path = self.den / summaries_path
        self.vectorstore_path = self.den / vectorstore_path
        self.download = download
        self.num_threads = num_threads

        self.database_path.mkdir(exist_ok=True, parents=True)
        self.summaries_path.mkdir(exist_ok=True, parents=True)

    # ---- abstract-ish methods ----
    def _search(self, query: str) -> list[dict[str, Any]]:
        raise NotImplementedError

    def _materialize(self, hit: dict[str, Any]) -> ItemMetadata:
        raise NotImplementedError

    def _id(self, hit_or_item: dict[str, Any]) -> str:
        raise NotImplementedError

    def _citation(self, item: ItemMetadata) -> str:
        # Subclass should format its ideal citation; fallback is ID or URL.
        return item.get("id") or item.get("url", "Unknown Source")

    # ---- optional hooks ----
    def _filter_hit(self, hit: dict[str, Any]) -> bool:
        return True

    def _postprocess_text(self, text: str, local_path: Optional[str]) -> str:
        # Default: optionally add image descriptions for PDFs
        if (
            self.process_images
            and local_path
            and local_path.lower().endswith(".pdf")
        ):
            try:
                descs = extract_and_describe_images(local_path)
                if any(descs):
                    text += "\n\n[Image Interpretations]\n" + "\n".join(descs)
            except Exception:
                pass
        return text

    # ---- shared nodes ----
    def _fetch_items(self, query: str) -> list[ItemMetadata]:
        hits = self._search(query)[: self.max_results] if self.download else []
        items: list[ItemMetadata] = []

        # If not downloading/scraping, try to load whatever is cached in database_path.
        if not self.download:
            for fname in os.listdir(self.database_path):
                if fname.lower().endswith((".pdf", ".txt", ".html")):
                    item_id = os.path.splitext(fname)[0]
                    local_path = os.path.join(self.database_path, fname)
                    full_text = ""
                    try:
                        if fname.lower().endswith(".pdf"):
                            full_text = read_pdf(local_path)
                        else:
                            with open(
                                local_path,
                                "r",
                                encoding="utf-8",
                                errors="ignore",
                            ) as f:
                                full_text = f.read()
                    except Exception as e:
                        full_text = f"[Error reading cached file: {e}]"
                    full_text = self._postprocess_text(full_text, local_path)
                    items.append({
                        "id": item_id,
                        "local_path": local_path,
                        "full_text": full_text,
                    })
            return items

        # Normal path: search → materialize each
        with ThreadPoolExecutor(
            max_workers=min(self.num_threads, max(1, len(hits)))
        ) as ex:
            futures = [
                ex.submit(self._materialize, h)
                for h in hits
                if self._filter_hit(h)
            ]
            for fut in as_completed(futures):
                try:
                    item = fut.result()
                    items.append(item)
                except Exception as e:
                    items.append({
                        "id": _hash(str(time.time())),
                        "full_text": f"[Error: {e}]",
                    })
        return items

    def _normalize_inputs(self, inputs) -> AcquisitionState:
        if isinstance(inputs, str):
            return AcquisitionState(context=inputs)
        elif isinstance(inputs, dict):
            return inputs
        else:
            raise TypeError(f"Invalid input for {self.__class__.__name__}")

    def format_result(self, state: AcquisitionState) -> str:
        if summary := state["final_summary"]:
            return summary

        # Fallback to dumping an empty string if `self.summarize=False`.
        # This only happens if a user disables summarization when configuring
        # URSA or if the final_summary is an empty string itself.
        return ""

    async def _search_query(self, state: AcquisitionState) -> AcquisitionState:
        """Generate a search query from the input search task (context)"""
        existing = state.get("query")
        if existing:
            return state

        context = state["context"]
        query = await self.llm.ainvoke(
            f"The user stated {context}. Generate between 1 and 8 words for a search query to address the users need. Return only the words to search."
        )
        state["query"] = query.content or context
        return state

    def _fetch_node(self, state: AcquisitionState) -> AcquisitionState:
        items = self._fetch_items(state["query"])
        return {**state, "items": items}

    def _summarize_node(self, state: AcquisitionState) -> AcquisitionState:
        prompt = ChatPromptTemplate.from_template("""
        You are an assistant responsible for summarizing retrieved content in the context of this task: {context}

        Summarize the content below:

        {retrieved_content}
        """)
        chain = prompt | self.llm | StrOutputParser()

        if "items" not in state or not state["items"]:
            return {**state, "summaries": None}

        summaries: list[Optional[str]] = [None] * len(state["items"])

        def process(i: int, item: ItemMetadata):
            item_id = item.get("id", f"item_{i}")
            out_path = os.path.join(
                self.summaries_path, f"{_safe_filename(item_id)}_summary.txt"
            )
            try:
                cleaned = remove_surrogates(item.get("full_text", ""))
                summary = chain.invoke(
                    {"retrieved_content": cleaned, "context": state["context"]},
                    config=self.build_config(tags=["acq", "summarize_each"]),
                )
            except Exception as e:
                summary = f"[Error summarizing item {item_id}: {e}]"
            with open(out_path, "w", encoding="utf-8") as f:
                f.write(summary)
            return i, summary

        with ThreadPoolExecutor(
            max_workers=min(self.num_threads, len(state["items"]))
        ) as ex:
            futures = [
                ex.submit(process, i, it) for i, it in enumerate(state["items"])
            ]
            for fut in as_completed(futures):
                i, s = fut.result()
                summaries[i] = s

        return {**state, "summaries": summaries}  # type: ignore

    def _rag_node(self, state: AcquisitionState) -> AcquisitionState:
        new_state = state.copy()
        rag_agent = RAGAgent(
            llm=self.llm,
            den=self.den,
            embedding=self.rag_embedding,
            vectorstore_path="rag_vectorstore",
            database_path=self.database_path.name,
        )
        new_state["final_summary"] = rag_agent.invoke(context=state["context"])[
            "summary"
        ]
        return new_state

    def _aggregate_node(self, state: AcquisitionState) -> AcquisitionState:
        if not state.get("summaries") or not state.get("items"):
            return {**state, "final_summary": None}

        blocks: list[str] = []
        for idx, (item, summ) in enumerate(
            zip(state["items"], state["summaries"])
        ):  # type: ignore
            cite = self._citation(item)
            blocks.append(f"[{idx + 1}] {cite}\n\nSummary:\n{summ}")

        combined = "\n\n" + ("\n\n" + "-" * 40 + "\n\n").join(blocks)
        with open(
            os.path.join(self.summaries_path, "summaries_combined.txt"),
            "w",
            encoding="utf-8",
        ) as f:
            f.write(combined)

        prompt = ChatPromptTemplate.from_template("""
        You are a scientific assistant extracting insights from multiple summaries.

        Here are the summaries:

        {Summaries}

        Your task is to read all the summaries and provide a response to this task: {context}
        """)
        chain = prompt | self.llm | StrOutputParser()

        final_summary = chain.invoke(
            {"Summaries": combined, "context": state["context"]},
            config=self.build_config(tags=["acq", "aggregate"]),
        )
        with open(
            os.path.join(self.summaries_path, "final_summary.txt"),
            "w",
            encoding="utf-8",
        ) as f:
            f.write(final_summary)

        return {**state, "final_summary": final_summary}

    def _build_graph(self):
        self.add_node(self._search_query)
        self.graph.set_entry_point("_search_query")
        self.add_node(self._fetch_node)
        self.graph.add_edge("_search_query", "_fetch_node")

        if self.summarize:
            if self.rag_embedding:
                self.add_node(self._rag_node)
                self.graph.add_edge("_fetch_node", "_rag_node")
                self.graph.set_finish_point("_rag_node")
            else:
                self.add_node(self._summarize_node)
                self.add_node(self._aggregate_node)

                self.graph.add_edge("_fetch_node", "_summarize_node")
                self.graph.add_edge("_summarize_node", "_aggregate_node")
                self.graph.set_finish_point("_aggregate_node")
        else:
            self.graph.set_finish_point("_fetch_node")

OSTIAgent

Bases: BaseAcquisitionAgent

Minimal OSTI.gov acquisition agent.

NOTE

• OSTI provides search endpoints that can return metadata including full-text links.
• Depending on your environment, you may prefer the public API or site scraping.
• Here we assume a JSON API that yields results with keys like: {'osti_id': '12345', 'title': '...', 'pdf_url': 'https://...pdf', 'landing_page': 'https://...'} Adapt field names if your OSTI integration differs.

Customize _search and _materialize to match your OSTI access path.

Source code in src/ursa/agents/acquisition_agents.py

class OSTIAgent(BaseAcquisitionAgent):
    """
    Minimal OSTI.gov acquisition agent.

    NOTE:
      - OSTI provides search endpoints that can return metadata including full-text links.
      - Depending on your environment, you may prefer the public API or site scraping.
      - Here we assume a JSON API that yields results with keys like:
            {'osti_id': '12345', 'title': '...', 'pdf_url': 'https://...pdf', 'landing_page': 'https://...'}
        Adapt field names if your OSTI integration differs.

    Customize `_search` and `_materialize` to match your OSTI access path.
    """

    def __init__(
        self,
        *args,
        api_base: str = "https://www.osti.gov/api/v1/records",
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.api_base = api_base

    def _id(self, hit_or_item: dict[str, Any]) -> str:
        if "osti_id" in hit_or_item:
            return str(hit_or_item["osti_id"])
        if "id" in hit_or_item:
            return str(hit_or_item["id"])
        if "landing_page" in hit_or_item:
            return _hash(hit_or_item["landing_page"])
        return _hash(json.dumps(hit_or_item))

    def _citation(self, item: ItemMetadata) -> str:
        t = item.get("title", "") or ""
        oid = item.get("id", "")
        return f"OSTI {oid}: {t}" if t else f"OSTI {oid}"

    def _search(self, query: str) -> list[dict[str, Any]]:
        """
        Adjust params to your OSTI setup. This call is intentionally simple;
        add paging/auth as needed.
        """
        params = {
            "q": query,
            "size": self.max_results,
        }
        try:
            r = requests.get(self.api_base, params=params, timeout=25)
            r.raise_for_status()
            data = r.json()
            # Normalize to a list of hits; adapt key if your API differs.
            if isinstance(data, dict) and "records" in data:
                hits = data["records"]
            elif isinstance(data, list):
                hits = data
            else:
                hits = []
            return hits[: self.max_results]
        except Exception as e:
            return [
                {
                    "id": _hash(query + str(time.time())),
                    "title": "Search error",
                    "error": str(e),
                }
            ]

    def _materialize(self, hit: dict[str, Any]) -> ItemMetadata:
        item_id = self._id(hit)
        title = hit.get("title") or hit.get("title_public", "") or ""
        landing = None
        local_path = ""
        full_text = ""

        try:
            pdf_url, landing_used, _ = resolve_pdf_from_osti_record(
                hit,
                headers={"User-Agent": "Mozilla/5.0"},
                unpaywall_email=os.environ.get("UNPAYWALL_EMAIL"),  # optional
            )

            if pdf_url:
                # Try to download as PDF (validate headers)
                with requests.get(
                    pdf_url,
                    headers={"User-Agent": "Mozilla/5.0"},
                    timeout=25,
                    allow_redirects=True,
                    stream=True,
                ) as r:
                    r.raise_for_status()
                    if _is_pdf_response(r):
                        fname = _derive_filename_from_cd_or_url(
                            r, f"osti_{item_id}.pdf"
                        )
                        local_path = os.path.join(self.database_path, fname)
                        _download_stream_to(local_path, r)
                        # Extract PDF text
                        try:
                            full_text = read_pdf(local_path)
                        except Exception as e:
                            full_text = (
                                f"[Downloaded but text extraction failed: {e}]"
                            )
                    else:
                        # Not a PDF; treat as HTML landing and parse text
                        landing = r.url
                        r.close()
            # If we still have no text, try scraping the DOE PAGES landing or citation page
            if not full_text:
                # Prefer DOE PAGES landing if present, else OSTI biblio
                landing = (
                    landing
                    or landing_used
                    or next(
                        (
                            link.get("href")
                            for link in hit.get("links", [])
                            if link.get("rel")
                            in ("citation_doe_pages", "citation")
                        ),
                        None,
                    )
                )
                if landing:
                    soup = _get_soup(
                        landing,
                        timeout=25,
                        headers={"User-Agent": "Mozilla/5.0"},
                    )
                    html_text = soup.get_text(" ", strip=True)
                    full_text = html_text[:1_000_000]  # keep it bounded
                    # Save raw HTML for cache/inspection
                    local_path = os.path.join(
                        self.database_path, f"{item_id}.html"
                    )
                    with open(local_path, "w", encoding="utf-8") as f:
                        f.write(str(soup))
                else:
                    full_text = "[No PDF or landing page text available.]"

        except Exception as e:
            full_text = f"[Error materializing OSTI {item_id}: {e}]"

        full_text = self._postprocess_text(full_text, local_path)
        return {
            "id": item_id,
            "title": title,
            "url": landing,
            "local_path": local_path,
            "full_text": full_text,
            "extra": {"raw_hit": hit},
        }

WebSearchAgent

Bases: BaseAcquisitionAgent

Uses DuckDuckGo Search (ddgs) to find pages, downloads HTML or PDFs, extracts text, and then follows the same summarize/RAG path.

Source code in src/ursa/agents/acquisition_agents.py

class WebSearchAgent(BaseAcquisitionAgent):
    """
    Uses DuckDuckGo Search (ddgs) to find pages, downloads HTML or PDFs,
    extracts text, and then follows the same summarize/RAG path.
    """

    def __init__(self, *args, user_agent: str = "Mozilla/5.0", **kwargs):
        super().__init__(*args, **kwargs)
        self.user_agent = user_agent
        if DDGS is None:
            raise ImportError(
                "duckduckgo-search (DDGS) is required for WebSearchAgentGeneric."
            )

    def _id(self, hit_or_item: dict[str, Any]) -> str:
        url = hit_or_item.get("href") or hit_or_item.get("url") or ""
        return (
            _hash(url)
            if url
            else hit_or_item.get("id", _hash(json.dumps(hit_or_item)))
        )

    def _citation(self, item: ItemMetadata) -> str:
        t = item.get("title", "") or ""
        u = item.get("url", "") or ""
        return f"{t} ({u})" if t else (u or item.get("id", "Web result"))

    def _search(self, query: str) -> list[dict[str, Any]]:
        results: list[dict[str, Any]] = []
        with DDGS() as ddgs:
            for r in ddgs.text(
                query, max_results=self.max_results, backend="auto"
            ):
                # r keys typically: title, href, body
                results.append(r)
        return results

    def _materialize(self, hit: dict[str, Any]) -> ItemMetadata:
        url = hit.get("href") or hit.get("url")
        title = hit.get("title", "")
        if not url:
            return {"id": self._id(hit), "title": title, "full_text": ""}

        headers = {"User-Agent": self.user_agent}
        local_path = ""
        full_text = ""
        item_id = self._id(hit)

        try:
            if _looks_like_pdf_url(url):
                local_path = os.path.join(
                    self.database_path, _safe_filename(item_id) + ".pdf"
                )
                _download(url, local_path)
                full_text = read_pdf(local_path)
            else:
                r = requests.get(url, headers=headers, timeout=20)
                r.raise_for_status()
                html = r.text
                local_path = os.path.join(
                    self.database_path, _safe_filename(item_id) + ".html"
                )
                with open(local_path, "w", encoding="utf-8") as f:
                    f.write(html)
                full_text = extract_main_text_only(html)
                # full_text = _basic_readable_text_from_html(html)
        except Exception as e:
            full_text = f"[Error retrieving {url}: {e}]"

        full_text = self._postprocess_text(full_text, local_path)
        return {
            "id": item_id,
            "title": title,
            "url": url,
            "local_path": local_path,
            "full_text": full_text,
            "extra": {"snippet": hit.get("body", "")},
        }

base

Base agent class providing telemetry, configuration, and execution abstractions.

This module defines the BaseAgent abstract class, which serves as the foundation for all agent implementations in the Ursa framework. It provides:

• Standardized initialization with LLM configuration
• Telemetry and metrics collection
• Thread and checkpoint management
• Input normalization and validation
• Execution flow control with invoke/stream methods
• Graph integration utilities for LangGraph compatibility
• Runtime enforcement of the agent interface contract

Agents built on this base class benefit from consistent behavior, observability, and integration capabilities while only needing to implement the core _invoke method.

AgentContext dataclass

Immutable context provided during graph execution

Source code in src/ursa/agents/base.py

@dataclass(frozen=True, kw_only=True)
class AgentContext:
    """Immutable context provided during graph execution"""

    llm: BaseChatModel
    """ Chat model for use during tool calls """

    workspace: Path
    """ Workspace path for the agent """

    agent_name: str
    """ Name for persisting the agent """

    group: str
    """ Name for the group for the agent to operate in. Controls data and endpoint availability"""

    den: Path
    """ Path to the persistent storage for the agent"""

    tool_character_limit: int = 30000
    """ Suggested limit on tool call responses """

    rag_tools: tuple[str, ...] = ()
    """Persisted URSA RAG collections bound to this agent as tools."""

    rag_tool_group: str | None = None
    """Group used to resolve persisted RAG tool collections."""

    rag_tool_return_k: int = 10
    """Number of chunks retrieved by persisted RAG tools."""

agent_name instance-attribute

Name for persisting the agent

den instance-attribute

Path to the persistent storage for the agent

group instance-attribute

Name for the group for the agent to operate in. Controls data and endpoint availability

llm instance-attribute

Chat model for use during tool calls

rag_tool_group = None class-attribute instance-attribute

Group used to resolve persisted RAG tool collections.

rag_tool_return_k = 10 class-attribute instance-attribute

Number of chunks retrieved by persisted RAG tools.

rag_tools = () class-attribute instance-attribute

Persisted URSA RAG collections bound to this agent as tools.

tool_character_limit = 30000 class-attribute instance-attribute

Suggested limit on tool call responses

workspace instance-attribute

Workspace path for the agent

AgentWithTools

Mixin that equips an agent with LangGraph tools management.

Source code in src/ursa/agents/base.py

class AgentWithTools:
    """Mixin that equips an agent with LangGraph tools management."""

    def __init__(
        self,
        *args,
        tools: list[BaseTool] | dict[str, BaseTool] | None = None,
        safe_codes: list[str] | None = None,
        handle_tool_errors=_default_tool_error_handler,
        **kwargs,
    ):
        self._tools: dict[str, BaseTool] = {}
        self.safe_codes = set(safe_codes or [])
        self.handle_tool_errors = handle_tool_errors
        self.tool_node = ToolNode([])
        self._apply_tools(tools, rebuild_graph=False)
        super().__init__(*args, **kwargs)
        self.tool_llm = self.llm.model_copy()

        if getattr(self, "rag_tools", ()):
            from ursa.rag.tools import build_rag_tools

            rag_tools = build_rag_tools(
                names=self.rag_tools,
                group=self.rag_tool_group or self.group or "default",
                llm=self.llm,
                embedding=getattr(self, "rag_tool_embedding", None),
                return_k=self.rag_tool_return_k,
            )
            if rag_tools:
                merged = dict(self._tools)
                merged.update({tool.name: tool for tool in rag_tools})
                self._apply_tools(merged, rebuild_graph=False)

    def __post_init__(self):
        super().__post_init__()

    def hook_storage_setup(self, store: BaseStore) -> None:
        if store is None:
            return
        for safe_code in self.safe_codes:
            store.put(
                ("workspace", "safe_codes"),
                safe_code,
                {},
            )

    async def ahook_storage_setup(self, store: BaseStore) -> None:
        if store is None:
            return
        for safe_code in self.safe_codes:
            await store.aput(
                ("workspace", "safe_codes"),
                safe_code,
                {},
            )

    @property
    def tools(self) -> dict[str, BaseTool]:
        return dict(self._tools)

    @tools.setter
    def tools(self, tools: dict[str, BaseTool] | list[BaseTool] | None):
        self._apply_tools(tools)

    def add_tool(self, tools: BaseTool | list[BaseTool]) -> None:
        bundle = tools if isinstance(tools, list) else [tools]
        merged = dict(self._tools)
        merged.update({tool.name: tool for tool in bundle})
        self._apply_tools(merged)

    async def add_mcp_tools(
        self,
        client: MultiServerMCPClient,
        tool_name: None | str | list[str] = None,
    ) -> None:
        """Add tools from an MCP client to the agent

        Args:
           client: the MCP client to add tools from
           tool_name: if provided, only add named tools
        """
        tools = await client.get_tools()
        if tool_name is not None:
            tool_name = (
                tool_name if isinstance(tool_name, list) else [tool_name]
            )
            tools = [tool for tool in tools if tool.name in tool_name]
        self.add_tool(tools)

    def remove_tool(self, tool_names: str | list[str]) -> None:
        names = tool_names if isinstance(tool_names, list) else [tool_names]
        trimmed = {
            name: tool
            for name, tool in self._tools.items()
            if name not in names
        }
        self._apply_tools(trimmed)

    def _apply_tools(
        self,
        tools: dict[str, BaseTool] | list[BaseTool] | None,
        *,
        rebuild_graph: bool = True,
    ) -> None:
        if tools is None:
            mapping: dict[str, BaseTool] = {}
        elif isinstance(tools, dict):
            mapping = dict(tools)
        else:
            mapping = {tool.name: tool for tool in tools}

        self._tools = mapping
        self.tool_node = ToolNode(
            list(self._tools.values()),
            handle_tool_errors=self.handle_tool_errors,
        )

        if rebuild_graph and hasattr(self, "build_graph"):
            self.__dict__.pop("compiled_graph", None)
            if hasattr(self, "graph"):
                self.build_graph()

add_mcp_tools(client, tool_name=None) async

Add tools from an MCP client to the agent

Parameters:

Name	Type	Description	Default
client	MultiServerMCPClient	the MCP client to add tools from	required
tool_name	None | str | list[str]	if provided, only add named tools	None

Source code in src/ursa/agents/base.py

async def add_mcp_tools(
    self,
    client: MultiServerMCPClient,
    tool_name: None | str | list[str] = None,
) -> None:
    """Add tools from an MCP client to the agent

    Args:
       client: the MCP client to add tools from
       tool_name: if provided, only add named tools
    """
    tools = await client.get_tools()
    if tool_name is not None:
        tool_name = (
            tool_name if isinstance(tool_name, list) else [tool_name]
        )
        tools = [tool for tool in tools if tool.name in tool_name]
    self.add_tool(tools)

BaseAgent

Bases: Generic[TState], ABC

Abstract base class for all agent implementations in the Ursa framework.

BaseAgent provides a standardized foundation for building LLM-powered agents with built-in telemetry, configuration management, and execution flow control. It handles common tasks like input normalization, thread management, metrics collection, and LangGraph integration.

Subclasses only need to implement the _invoke method to define their core functionality, while inheriting standardized invocation patterns, telemetry, and graph integration capabilities. The class enforces a consistent interface through runtime checks that prevent subclasses from overriding critical methods like invoke().

The agent supports both direct invocation with inputs and streaming responses, with automatic tracking of token usage, execution time, and other metrics. It also provides utilities for integrating with LangGraph through node wrapping and configuration.

Subclass Inheritance Guidelines

• Must Override: _invoke() - Define your agent's core functionality
• Can Override: _stream() - Enable streaming support _normalize_inputs() - Customize input handling Various helper methods (_default_node_tags, etc.)
• Never Override: invoke() - Final method with runtime enforcement stream() - Handles telemetry and delegates to _stream call() - Delegates to invoke Other public methods (build_config, write_state, add_node)

To create a custom agent, inherit from this class and implement the _invoke method:

class MyAgent(BaseAgent):
    def _invoke(self, inputs: Mapping[str, Any], **config: Any) -> Any:
        # Process inputs and return results
        ...

Source code in src/ursa/agents/base.py

class BaseAgent(Generic[TState], ABC):
    """Abstract base class for all agent implementations in the Ursa framework.

    BaseAgent provides a standardized foundation for building LLM-powered agents with
    built-in telemetry, configuration management, and execution flow control. It handles
    common tasks like input normalization, thread management, metrics collection, and
    LangGraph integration.

    Subclasses only need to implement the _invoke method to define their core
    functionality, while inheriting standardized invocation patterns, telemetry, and
    graph integration capabilities. The class enforces a consistent interface through
    runtime checks that prevent subclasses from overriding critical methods like
    invoke().

    The agent supports both direct invocation with inputs and streaming responses, with
    automatic tracking of token usage, execution time, and other metrics. It also
    provides utilities for integrating with LangGraph through node wrapping and
    configuration.

    Subclass Inheritance Guidelines:
        - Must Override: _invoke() - Define your agent's core functionality
        - Can Override: _stream() - Enable streaming support
                        _normalize_inputs() - Customize input handling
                        Various helper methods (_default_node_tags, etc.)
        - Never Override: invoke() - Final method with runtime enforcement
                          stream() - Handles telemetry and delegates to _stream
                          __call__() - Delegates to invoke
                          Other public methods (build_config, write_state, add_node)

    To create a custom agent, inherit from this class and implement the _invoke method:

    ```python
    class MyAgent(BaseAgent):
        def _invoke(self, inputs: Mapping[str, Any], **config: Any) -> Any:
            # Process inputs and return results
            ...
    ```
    """

    # This will be shared across all BaseAgent instances.
    _invoke_depth: int = 0

    _TELEMETRY_KW = {
        "raw_debug",
        "save_json",
        "save_otel",
        "metrics_path",
        "save_raw_snapshot",
        "save_raw_records",
        "otel_endpoint",
        "otel_headers",
    }

    _CONTROL_KW = {"config", "recursion_limit", "tags", "metadata", "callbacks"}

    state_type: type[TState] = dict

    def __init__(
        self,
        llm: BaseChatModel,
        workspace: Optional[Path] = None,
        agent_name: Optional[str] = None,
        group: Optional[str] = "default",
        checkpointer: Optional[BaseCheckpointSaver] = None,
        enable_metrics: bool = True,
        metrics_dir: str = "ursa_metrics",  # dir to save metrics, with a default
        autosave_metrics: bool = True,
        otel_metrics: bool = False,
        thread_id: Optional[str] = None,
        tokens_before_summarize: int = 50000,
        messages_to_keep: int = 20,
        max_single_tool_message_tokens: int = 100000,
        rag_tools: str | Sequence[str] | None = None,
        rag_tool_group: str | None = None,
        rag_tool_embedding: Embeddings | None = None,
        rag_tool_return_k: int = 10,
    ):
        """Initializes the base agent with a language model and optional configurations.

        Args:
            llm: a BaseChatModel instance.
            checkpointer: Optional checkpoint saver for persisting agent state.
            enable_metrics: Whether to collect performance and usage metrics.
            metrics_dir: Directory path where metrics will be saved.
            autosave_metrics: Whether to automatically save metrics to disk.
            thread_id: Unique identifier for this agent instance. Generated if not
                       provided.
        """
        self.llm: BaseChatModel = llm
        self.workspace = Path(workspace or ".")
        self.agent_name = agent_name
        self.group = validate_group_name(group)
        self.tokens_before_summarize = tokens_before_summarize
        self.messages_to_keep = messages_to_keep
        self.max_single_tool_message_tokens = max_single_tool_message_tokens
        from ursa.rag.persistence import normalize_rag_tool_names

        self.rag_tools = tuple(normalize_rag_tool_names(rag_tools))
        self.rag_tool_group = rag_tool_group
        self.rag_tool_embedding = rag_tool_embedding
        self.rag_tool_return_k = rag_tool_return_k
        enforce_model_group_policy(self.llm, self.group)
        if (
            not group_root_dir(self.group).exists()
            and self.group != DEFAULT_GROUP_NAME
        ):
            raise ValueError(
                (
                    f"Group '{self.group}' does not exist. "
                    f"Please use `ursa create-group {self.group} <group_config_file>` to create"
                )
            )
        set_checkpointer = True if checkpointer else False
        set_name = True if agent_name else False
        self.checkpointer = checkpointer
        self._checkpointer_was_provided = set_checkpointer
        self._async_checkpointer: BaseCheckpointSaver | None = None
        self._async_storage: BaseStore | None = None
        self._async_compiled_graph: CompiledStateGraph | None = None
        persist_agent = (agent_name is not None) or set_checkpointer
        if persist_agent:
            if set_checkpointer:
                if set_name:
                    print(
                        (
                            "[WARNING]: Both checkpointer and den persistence set."
                            " Using checkpointer, but only use one in the future."
                        )
                    )
                self.den = self.workspace
            else:
                den_name = Path(self.group) / "agents" / agent_name
                self.den = group_agents_dir(self.group) / agent_name
                if not self.den.exists():
                    print(f"[Agent Created]: {den_name}")
                self.checkpointer = Checkpointer.from_workspace(self.den)
        else:
            # Keep current behavior if the user is not persisting.
            self.den = self.workspace
        self.thread_id = thread_id or "ursa"
        self.telemetry = Telemetry(
            enable=enable_metrics,
            output_dir=self.den.joinpath(metrics_dir),
            save_json_default=autosave_metrics,
        )

        self.workspace.mkdir(exist_ok=True, parents=True)
        self.den.mkdir(exist_ok=True, parents=True)

    @property
    def name(self) -> str:
        """Agent name."""
        return self.__class__.__name__

    @property
    def context(self) -> AgentContext:
        """Immutable run-scoped information provided to the Agent's graph"""
        return AgentContext(
            llm=self.llm,
            workspace=self.workspace,
            agent_name=self.agent_name,
            group=self.group,
            den=self.den,
            rag_tools=self.rag_tools,
            rag_tool_group=self.rag_tool_group,
            rag_tool_return_k=self.rag_tool_return_k,
        )

    @staticmethod
    def _dedupe_callbacks(callbacks: list[Any]) -> list[Any]:
        """Preserve callback order while removing duplicate handler instances."""
        deduped: list[Any] = []
        seen: set[int] = set()
        for callback in callbacks:
            marker = id(callback)
            if marker in seen:
                continue
            deduped.append(callback)
            seen.add(marker)
        return deduped

    def events(self, config: RunnableConfig | None = None) -> AgentEvents:
        """Return a helper for emitting structured agent progress events."""
        return AgentEvents(agent=self.name, config=config)

    def add_node(
        self,
        f: Callable[..., Mapping[str, Any]],
        node_name: Optional[str] = None,
        agent_name: Optional[str] = None,
        **kwargs,
    ) -> StateGraph:
        """Add a node to the state graph with token usage tracking.

        This method adds a function as a node to the state graph, wrapping it to track
        token usage during execution. The node is identified by either the provided
        node_name or the function's name.

        Args:
            f: The function to add as a node. Should return a mapping of string keys to
                any values.
            node_name: Optional name for the node. If not provided, the function's name
                will be used.
            agent_name: Optional agent name for tracking. If not provided, the agent's
                name in snake_case will be used.

        Returns:
            The updated StateGraph with the new node added.
        """
        _node_name = node_name or f.__name__
        _agent_name = agent_name or _to_snake(self.name)
        wrapped_node = self._wrap_node(f, _node_name, _agent_name)
        return self.graph.add_node(_node_name, wrapped_node, **kwargs)

    def write_state(self, filename: str, state: dict) -> None:
        """Writes agent state to a JSON file.

        Serializes the provided state dictionary to JSON format and writes it to the
        specified file. The JSON is written with non-ASCII characters preserved.

        Args:
            filename: Path to the file where state will be written.
            state: Dictionary containing the agent state to be serialized.
        """
        json_state = dumps(state, ensure_ascii=False)
        with open(filename, "w") as f:
            f.write(json_state)

    def build_config(self, **overrides) -> dict:
        """Constructs a config dictionary for agent operations with telemetry support.

        This method creates a standardized configuration dictionary that includes thread
        identification, telemetry callbacks, and other metadata needed for agent
        operations. The configuration can be customized through override parameters.

        Args:
            **overrides: Optional configuration overrides that can include keys like
                'recursion_limit', 'configurable', 'metadata', 'tags', etc.

        Returns:
            dict: A complete configuration dictionary with all necessary parameters.
        """
        # Create the base configuration with essential fields.
        base = {
            "configurable": {"thread_id": self.thread_id},
            "metadata": {
                "thread_id": self.thread_id,
                "telemetry_run_id": self.telemetry.context.get("run_id"),
            },
            "tags": [self.name],
            "callbacks": [
                *self.telemetry.callbacks,
                DEFAULT_EVENT_LOGGING_HANDLER,
            ],
        }

        # Try to determine the model name from either direct or nested attributes
        model_name = getattr(self, "llm_model", None) or getattr(
            getattr(self, "llm", None), "model", None
        )

        # Add model name to metadata if available
        if model_name:
            base["metadata"]["model"] = model_name

        # Handle configurable dictionary overrides by merging with base configurable
        if "configurable" in overrides and isinstance(
            overrides["configurable"], dict
        ):
            base["configurable"].update(overrides.pop("configurable"))

        # Handle metadata dictionary overrides by merging with base metadata
        if "metadata" in overrides and isinstance(overrides["metadata"], dict):
            base["metadata"].update(overrides.pop("metadata"))

        # Merge tags from caller-provided overrides, avoid duplicates
        if "tags" in overrides and isinstance(overrides["tags"], list):
            base["tags"] = base["tags"] + [
                t for t in overrides.pop("tags") if t not in base["tags"]
            ]

        if "callbacks" in overrides:
            callbacks = merge_configs(
                {"callbacks": base["callbacks"]},
                {"callbacks": overrides.pop("callbacks")},
            ).get("callbacks")
            if isinstance(callbacks, list):
                callbacks = self._dedupe_callbacks(callbacks)
            base["callbacks"] = callbacks

        # Apply any remaining overrides directly to the base configuration
        base.update(overrides)

        return base

    def _invoke_engine(
        self,
        invoke_method,
        inputs: Optional[InputLike] = None,
        raw_debug: bool = False,
        save_json: Optional[bool] = None,
        save_otel: Optional[bool] = None,
        metrics_path: Optional[str] = None,
        otel_endpoint: Optional[str] = None,
        otel_headers: Optional[str] = None,
        save_raw_snapshot: Optional[bool] = None,
        save_raw_records: Optional[bool] = None,
        config: Optional[dict] = None,
        **kwargs: Any,
    ):
        BaseAgent._invoke_depth += 1
        invoke_config = dict(config or {})

        try:
            # Start telemetry tracking for the top-level invocation
            if BaseAgent._invoke_depth == 1:
                self.telemetry.begin_run(
                    agent=self.name, thread_id=self.thread_id
                )

            # Handle the case where inputs are provided as keyword arguments
            if inputs is None:
                # Separate kwargs into input parameters and control parameters
                kw_inputs: dict[str, Any] = {}
                control_kwargs: dict[str, Any] = {}
                for k, v in kwargs.items():
                    if k in self._TELEMETRY_KW or k in self._CONTROL_KW:
                        control_kwargs[k] = v
                    else:
                        kw_inputs[k] = v
                inputs = kw_inputs

                # Only control kwargs remain for further processing
                kwargs = control_kwargs

            # Handle the case where inputs are provided as a positional argument
            else:
                # Ensure no ambiguous keyword arguments are present
                for k in kwargs.keys():
                    if not (k in self._TELEMETRY_KW or k in self._CONTROL_KW):
                        raise TypeError(
                            f"Unexpected keyword argument '{k}'. "
                            "Pass inputs as a single mapping or omit the positional "
                            "inputs and pass them as keyword arguments."
                        )

            # Allow subclasses to normalize or transform the input format
            normalized = self._normalize_inputs(inputs)

            # Delegate to the subclass implementation with the normalized inputs
            # and any control parameters
            return invoke_method(normalized, **invoke_config, **kwargs)

        finally:
            # Clean up the invocation depth tracking
            BaseAgent._invoke_depth -= 1

            # For the top-level invocation, finalize telemetry and generate outputs
            if BaseAgent._invoke_depth == 0:
                self.telemetry.render(
                    raw=raw_debug,
                    save_json=save_json,
                    save_otel=save_otel,
                    filepath=metrics_path,
                    otel_endpoint=otel_endpoint,
                    otel_headers=otel_headers,
                    save_raw_snapshot=save_raw_snapshot,
                    save_raw_records=save_raw_records,
                )

    async def _ainvoke_engine(
        self,
        invoke_method,
        inputs: Optional[InputLike] = None,
        raw_debug: bool = False,
        save_json: Optional[bool] = None,
        save_otel: Optional[bool] = None,
        metrics_path: Optional[str] = None,
        otel_endpoint: Optional[str] = None,
        otel_headers: Optional[str] = None,
        save_raw_snapshot: Optional[bool] = None,
        save_raw_records: Optional[bool] = None,
        config: Optional[dict] = None,
        **kwargs: Any,
    ):
        BaseAgent._invoke_depth += 1
        invoke_config = dict(config or {})

        try:
            if BaseAgent._invoke_depth == 1:
                self.telemetry.begin_run(
                    agent=self.name, thread_id=self.thread_id
                )

            if inputs is None:
                kw_inputs: dict[str, Any] = {}
                control_kwargs: dict[str, Any] = {}
                for k, v in kwargs.items():
                    if k in self._TELEMETRY_KW or k in self._CONTROL_KW:
                        control_kwargs[k] = v
                    else:
                        kw_inputs[k] = v
                inputs = kw_inputs
                kwargs = control_kwargs
            else:
                for k in kwargs.keys():
                    if not (k in self._TELEMETRY_KW or k in self._CONTROL_KW):
                        raise TypeError(
                            f"Unexpected keyword argument '{k}'. "
                            "Pass inputs as a single mapping or omit the "
                            "positional inputs and pass them as keyword "
                            "arguments."
                        )

            normalized = self._normalize_inputs(inputs)
            return await invoke_method(normalized, **invoke_config, **kwargs)

        finally:
            BaseAgent._invoke_depth -= 1

            if BaseAgent._invoke_depth == 0:
                self.telemetry.render(
                    raw=raw_debug,
                    save_json=save_json,
                    save_otel=save_otel,
                    filepath=metrics_path,
                    otel_endpoint=otel_endpoint,
                    otel_headers=otel_headers,
                    save_raw_snapshot=save_raw_snapshot,
                    save_raw_records=save_raw_records,
                )

    # NOTE: The `invoke` method uses the PEP 570 `/,*` notation to explicitly state which
    # arguments can and cannot be passed as positional or keyword arguments.
    @final
    def invoke(
        self,
        inputs: Optional[InputLike] = None,
        /,
        *,
        raw_debug: bool = False,
        save_json: Optional[bool] = None,
        save_otel: Optional[bool] = None,
        metrics_path: Optional[str] = None,
        save_raw_snapshot: Optional[bool] = None,
        save_raw_records: Optional[bool] = None,
        config: Optional[dict] = None,
        **kwargs: Any,
    ) -> Any:
        """Executes the agent with the provided inputs and configuration.

        This is the main entry point for agent execution. It handles input normalization,
        telemetry tracking, and proper execution context management. The method supports
        flexible input formats - either as a positional argument or as keyword arguments.

        Args:
            inputs: Optional positional input to the agent. If provided, all non-control
                keyword arguments will be rejected to avoid ambiguity.
            raw_debug: If True, displays raw telemetry data for debugging purposes.
            save_json: If True, saves telemetry data as JSON.
            save_otel: If True, saves telemetry data to OpenTelemetry endpoint.
            metrics_path: Optional file path where telemetry metrics should be saved.
            save_raw_snapshot: If True, saves a raw snapshot of the telemetry data.
            save_raw_records: If True, saves raw telemetry records.
            config: Optional configuration dictionary to override default settings.
            **kwargs: Additional keyword arguments that can be either:
                - Input parameters (when no positional input is provided)
                - Control parameters recognized by the agent

        Returns:
            The result of the agent's execution.

        Raises:
            TypeError: If both positional inputs and non-control keyword arguments are
                provided simultaneously.
        """
        return self._invoke_engine(
            invoke_method=self._invoke,
            inputs=inputs,
            raw_debug=raw_debug,
            save_json=save_json,
            save_otel=save_otel,
            metrics_path=metrics_path,
            save_raw_snapshot=save_raw_snapshot,
            save_raw_records=save_raw_records,
            config=config,
            **kwargs,
        )

    # NOTE: The `ainvoke` method uses the PEP 570 `/,*` notation to explicitly state which
    # arguments can and cannot be passed as positional or keyword arguments.
    @final
    async def ainvoke(
        self,
        inputs: Optional[InputLike] = None,
        /,
        *,
        raw_debug: bool = False,
        save_json: Optional[bool] = None,
        save_otel: Optional[bool] = None,
        metrics_path: Optional[str] = None,
        save_raw_snapshot: Optional[bool] = None,
        save_raw_records: Optional[bool] = None,
        config: Optional[dict] = None,
        **kwargs: Any,
    ) -> Any:
        """Asynchrnously executes the agent with the provided inputs and configuration.

        (Async version of `invoke`.)

        This is the main entry point for agent execution. It handles input normalization,
        telemetry tracking, and proper execution context management. The method supports
        flexible input formats - either as a positional argument or as keyword arguments.

        Args:
            inputs: Optional positional input to the agent. If provided, all non-control
                keyword arguments will be rejected to avoid ambiguity.
            raw_debug: If True, displays raw telemetry data for debugging purposes.
            save_json: If True, saves telemetry data as JSON.
            save_otel: If True, saves telemetry data to OpenTelemetry endpoint.
            metrics_path: Optional file path where telemetry metrics should be saved.
            save_raw_snapshot: If True, saves a raw snapshot of the telemetry data.
            save_raw_records: If True, saves raw telemetry records.
            config: Optional configuration dictionary to override default settings.
            **kwargs: Additional keyword arguments that can be either:
                - Input parameters (when no positional input is provided)
                - Control parameters recognized by the agent

        Returns:
            The result of the agent's execution.

        Raises:
            TypeError: If both positional inputs and non-control keyword arguments are
                provided simultaneously.
        """
        return await self._ainvoke_engine(
            invoke_method=self._ainvoke,
            inputs=inputs,
            raw_debug=raw_debug,
            save_json=save_json,
            save_otel=save_otel,
            metrics_path=metrics_path,
            save_raw_snapshot=save_raw_snapshot,
            save_raw_records=save_raw_records,
            config=config,
            **kwargs,
        )

    def format_query(self, prompt: str, state: TState | None = None) -> TState:
        """Format a plain text prompt into the agent's input schema
        possibly incorporating the prior state.

        Agents should override this method for their operation
        """

        if state is not None and "messages" in state:
            state["messages"].append(HumanMessage(content=str(prompt)))
            return state
        return self._normalize_inputs(prompt)

    def format_result(self, result: TState) -> str:
        """Extracts a plain text response from the Agent's output schema

        Agents should override this method for their operation
        """

        if "messages" in result:
            if isinstance(result["messages"], list) and isinstance(
                result["messages"][-1], BaseMessage
            ):
                return result["messages"][-1].text
        raise NotImplementedError()

    def _message_tool_call_ids(self, msg: Any) -> list[str]:
        """Return LangChain tool-call IDs requested by a message, if any."""
        return [
            call["id"]
            for call in getattr(msg, "tool_calls", []) or []
            if "id" in call
        ]

    def _patch_dangling(
        self,
        state: Mapping[str, Any],
        summarized: bool = False,
        config: RunnableConfig | None = None,
    ) -> tuple[Mapping[str, Any], bool]:
        """Patch missing tool responses in a LangGraph message state.

        Tool-calling chat models require every AI tool call to be followed by a
        corresponding ``ToolMessage``. If a tool result is missing because a run
        timed out, failed, or summarization trimmed it away, later model calls can
        fail with provider-side message validation errors. This utility inserts a
        synthetic response for any dangling tool call ID and truncates extremely
        large tool messages.

        Args:
            state: A graph state containing a ``messages`` list.
            summarized: Existing overwrite flag from prior state mutation.

        Returns:
            ``(new_state, full_overwrite_required)``. When the boolean is true,
            callers using an ``add_messages`` reducer should return an
            ``Overwrite(new_state["messages"])`` update rather than appending.
        """
        if "messages" not in state:
            return state, summarized

        new_state = deepcopy(state)
        events = self.events(config)
        dangling_response = (
            "Response Not Found from tool. "
            "May have timed out or been forgotten due to summarization."
        )

        pending_tool_ids: list[str] = []
        for msg in new_state["messages"]:
            if isinstance(msg, ToolMessage):
                if (
                    self.max_single_tool_message_tokens is not None
                    and count_tokens_approximately([msg])
                    > self.max_single_tool_message_tokens
                ):
                    msg.content = "Message too long - truncated."
                    summarized = True
                try:
                    pending_tool_ids.remove(msg.tool_call_id)
                except ValueError:
                    # Orphan tool messages are not ideal, but do not create the
                    # provider-side validation error that missing tool responses do.
                    pass
                continue

            pending_tool_ids.extend(self._message_tool_call_ids(msg))

        if pending_tool_ids:
            summarized = True
            events.emit(
                "Dangling tool calls patched",
                stage="dangling_tool_calls",
                tool_call_ids=list(pending_tool_ids),
            )
            for tool_id in pending_tool_ids:
                for msg_ind, msg in enumerate(new_state["messages"]):
                    if tool_id in self._message_tool_call_ids(msg):
                        new_state["messages"].insert(
                            msg_ind + 1,
                            ToolMessage(
                                content=dangling_response,
                                tool_call_id=tool_id,
                            ),
                        )
                        break

        return new_state, summarized

    def _summarize_context(
        self,
        state: Mapping[str, Any],
        *,
        system_prompt: str | None = None,
        summary_prompt: str | None = None,
        config: RunnableConfig | None = None,
    ) -> tuple[Mapping[str, Any], bool]:
        """Summarize long ``state['messages']`` histories.

        This preserves the first/system message and the last
        ``self.messages_to_keep`` messages, replacing the middle of the
        transcript with an LLM-generated summary when the approximate token count
        exceeds ``self.tokens_before_summarize``.

        The method also avoids creating dangling tool-call pairs when possible by
        moving tool responses that correspond to summarized-away calls into the
        summarized block.
        """
        if "messages" not in state:
            return state, False

        new_state = deepcopy(state)
        events = self.events(config)
        messages = list(new_state.get("messages") or [])
        if len(messages) <= 1:
            return new_state, False

        tokens_before = count_tokens_approximately(messages[1:])
        if tokens_before <= self.tokens_before_summarize:
            return new_state, False

        events.emit(
            message=f"Summarizing ~{tokens_before} token history to compress context",
            stage="summarize_context",
        )
        keep_count = max(0, int(self.messages_to_keep or 0))
        body_messages = messages[1:]
        if keep_count == 0:
            conversation_to_summarize = body_messages
            conversation_to_keep = []
        elif len(body_messages) > keep_count:
            conversation_to_summarize = body_messages[:-keep_count]
            conversation_to_keep = body_messages[-keep_count:]
        else:
            # Nothing can be summarized while still preserving the requested tail.
            return new_state, False

        tool_ids: list[str] = []
        for msg in conversation_to_summarize:
            tool_ids.extend(self._message_tool_call_ids(msg))
            if isinstance(msg, ToolMessage):
                try:
                    tool_ids.remove(msg.tool_call_id)
                except ValueError:
                    pass

        if tool_ids:
            events.emit(
                "Preserving tool responses during summarization",
                stage="summarize_context",
                tool_call_ids=list(tool_ids),
            )
            keep_copy = list(conversation_to_keep)
            for msg in keep_copy:
                if (
                    isinstance(msg, ToolMessage)
                    and msg.tool_call_id in tool_ids
                ):
                    conversation_to_summarize.append(msg)
                    conversation_to_keep.remove(msg)
                    tool_ids.remove(msg.tool_call_id)

        if tool_ids:
            events.emit(
                "Dangling tool calls found during summarization",
                stage="summarize_context",
                tool_call_ids=list(tool_ids),
            )

        summarize_system_message = SystemMessage(
            content="""
        Your only tasks is to provide a detailed, comprehensive summary of the following
        conversation.

        Your summary will be the only information retained from the conversation, so ensure
        it contains all details that need to be remembered to meet the goals of the work.

        The conversation to summarize is:
        """
        )
        to_summarize = [summarize_system_message] + conversation_to_summarize
        to_summarize += [
            HumanMessage(
                content="Summarize that conversation per your instruction."
            )
        ]
        summary = self.tool_llm.invoke(to_summarize)

        first_message = messages[0]
        if system_prompt is not None:
            first_message = SystemMessage(content=system_prompt)

        new_state["messages"] = [first_message, summary] + conversation_to_keep
        return new_state, True

    def prepare_messages_context(
        self,
        state: Mapping[str, Any],
        *,
        system_prompt: str | None = None,
        summary_prompt: str | None = None,
        patch_dangling: bool = True,
        summarize: bool = True,
    ) -> tuple[Mapping[str, Any], bool]:
        """Apply standard message-history maintenance before an LLM call.

        Returns ``(new_state, full_overwrite_required)``. Graph nodes using the
        ``add_messages`` reducer should use :meth:`messages_update` with the
        returned flag to avoid appending duplicate history after summarization or
        dangling-tool patching.
        """
        new_state = deepcopy(state)
        full_overwrite = False
        if summarize:
            new_state, full_overwrite = self._summarize_context(
                new_state,
                system_prompt=system_prompt,
                summary_prompt=summary_prompt,
            )
        if patch_dangling:
            new_state, full_overwrite = self._patch_dangling(
                new_state, full_overwrite
            )
        return new_state, full_overwrite

    def messages_update(
        self,
        state: Mapping[str, Any],
        message_update: Any,
        *,
        full_overwrite: bool = False,
        extra: Mapping[str, Any] | None = None,
    ) -> dict[str, Any]:
        """Build a LangGraph-safe update for states with a ``messages`` reducer."""
        update = dict(extra or {})
        if full_overwrite:
            messages = list(state.get("messages") or [])
            if message_update is None:
                update["messages"] = Overwrite(messages)
            else:
                if isinstance(message_update, list):
                    messages.extend(message_update)
                else:
                    messages.append(message_update)
                update["messages"] = Overwrite(messages)
        else:
            update["messages"] = message_update
        return update

    def _normalize_inputs(self, inputs: InputLike) -> Mapping[str, Any]:
        """Normalizes various input formats into a standardized mapping.

        This method converts different input types into a consistent dictionary format
        that can be processed by the agent. String inputs are wrapped as messages, while
        mappings are passed through unchanged.

        Args:
            inputs: The input to normalize. Can be a string (which will be converted to a
                message) or a mapping (which will be returned as-is).

        Returns:
            A mapping containing the normalized inputs, with keys appropriate for agent
            processing.

        Raises:
            TypeError: If the input type is not supported (neither string nor mapping).
        """
        if isinstance(inputs, str):
            # Adjust to your message type
            return {"messages": [HumanMessage(content=inputs)]}
        if isinstance(inputs, Mapping):
            return inputs
        raise TypeError(f"Unsupported input type: {type(inputs)}")

    @cached_property
    def compiled_graph(self) -> CompiledStateGraph:
        """Return the sync compiled StateGraph application for the agent."""
        graph = self.build_graph()
        compiled = graph.compile(
            checkpointer=self.checkpointer,
            store=self.storage,
        ).with_config({"recursion_limit": 50000})
        return self._finalize_graph(compiled)

    async def _aget_async_compiled_graph(self) -> CompiledStateGraph:
        """Return an async-compatible compiled graph for the agent.

        LangGraph SQLite checkpointers are not interchangeable between sync
        and async execution. A graph compiled with ``SqliteSaver`` cannot be
        awaited with ``ainvoke``. Build a separate graph using async SQLite
        persistence resources for async runs.
        """
        if self._async_compiled_graph is None:
            graph = self.build_graph()
            compiled = graph.compile(
                checkpointer=await self._aget_async_checkpointer(),
                store=await self._aget_async_storage(),
            ).with_config({"recursion_limit": 50000})
            self._async_compiled_graph = self._finalize_graph(compiled)
        return self._async_compiled_graph

    async def _aget_async_checkpointer(self) -> BaseCheckpointSaver | None:
        """Return an async-compatible checkpointer for async graph runs."""
        if self.checkpointer is None:
            return None

        if isinstance(self.checkpointer, AsyncSqliteSaver):
            return self.checkpointer

        if self._async_checkpointer is not None:
            return self._async_checkpointer

        if isinstance(self.checkpointer, SqliteSaver):
            if self._checkpointer_was_provided:
                raise RuntimeError(
                    "This agent was configured with a synchronous SqliteSaver "
                    "checkpointer. Async agent execution requires an async "
                    "checkpointer, such as AsyncSqliteSaver, or URSA-managed "
                    "persistence via agent_name."
                )
            self._async_checkpointer = await Checkpointer.async_from_workspace(
                self.den
            )
            return self._async_checkpointer

        if self._checkpointer_has_async_methods(self.checkpointer):
            return self.checkpointer

        raise RuntimeError(
            "The configured checkpointer does not appear to support LangGraph "
            "async checkpoint methods. Use an async checkpointer for "
            "agent.ainvoke(...), or use URSA-managed persistence via "
            "agent_name."
        )

    @staticmethod
    def _checkpointer_has_async_methods(checkpointer: Any) -> bool:
        required = ("aget_tuple", "alist", "aput", "aput_writes")
        return all(
            inspect.iscoroutinefunction(getattr(type(checkpointer), name, None))
            or inspect.isasyncgenfunction(
                getattr(type(checkpointer), name, None)
            )
            for name in required
        )

    @cached_property
    def storage(self) -> BaseStore:
        """Create a sync SQLite-backed LangGraph store."""
        store_path = self.den / "graph_store.sqlite"
        conn = sqlite3.connect(
            store_path, check_same_thread=False, isolation_level=None
        )
        store = SqliteStore(conn)
        store.setup()
        self.hook_storage_setup(store)
        return store

    async def _aget_async_storage(self) -> BaseStore:
        """Create an async SQLite-backed LangGraph store."""
        if self._async_storage is None:
            import aiosqlite

            store_path = self.den / "graph_store.sqlite"
            conn = await aiosqlite.connect(store_path, isolation_level=None)
            store = AsyncSqliteStore(conn)
            await store.setup()
            await self.ahook_storage_setup(store)
            self._async_storage = store
        return self._async_storage

    def hook_storage_setup(self, store: BaseStore) -> None:
        pass

    async def ahook_storage_setup(self, store: BaseStore) -> None:
        pass

    def close(self) -> None:
        """Close sync SQLite resources owned by this agent when possible."""
        if "storage" in self.__dict__:
            self.storage.conn.close()
        if isinstance(self.checkpointer, SqliteSaver):
            self.checkpointer.conn.close()

    async def aclose(self) -> None:
        """Close async SQLite resources owned by this agent when possible."""
        if self._async_storage is not None:
            await self._async_storage.__aexit__(None, None, None)
            await self._async_storage.conn.close()
            self._async_storage = None
            self._async_compiled_graph = None
        if isinstance(self._async_checkpointer, AsyncSqliteSaver):
            await self._async_checkpointer.conn.close()
            self._async_checkpointer = None
            self._async_compiled_graph = None

    @final
    def build_graph(self) -> StateGraph:
        """Build and return the StateGraph backing this agent."""
        self.graph = StateGraph(
            self.state_type,
            context_schema=AgentContext,
        )
        self._build_graph()
        return self.graph

    @abstractmethod
    def _build_graph(self) -> None:
        """Construct the StateGraph for this agent without compiling.

        Called during `__post_init__()` after the Agent has been fully
        Initialized (`__post_init__` is called after `__init__`) to
        instantiate `self.graph`

        Agents should implement this to define their their behavior.

        Agents should treat `self.graph` as read-only
        """
        ...

    def _finalize_graph(
        self, graph_app: CompiledStateGraph
    ) -> CompiledStateGraph:
        """Hook for subclasses to wrap or modify the compiled graph."""
        return graph_app

    def _tool_is_async_only(self, tool: Any) -> bool:
        """Return True for tools that can only be invoked asynchronously.

        MCP tools are commonly exposed as StructuredTool instances with a
        coroutine implementation but no synchronous function implementation.
        Those raise errors like:

            "StructuredTool does not support sync invocation."

        when called via `.invoke()`.
        """

        func = getattr(tool, "func", None)
        coroutine = getattr(tool, "coroutine", None)
        return func is None and coroutine is not None

    def _has_async_only_tools(self) -> bool:
        tools_obj = getattr(self, "tools", None)
        if not tools_obj:
            return False

        try:
            tool_iter = (
                tools_obj.values() if isinstance(tools_obj, dict) else tools_obj
            )
        except Exception:
            return False

        return any(self._tool_is_async_only(t) for t in tool_iter)

    def _invoke(self, input, **config):
        config = self.build_config(**config)

        # If we have async-only tools (e.g. MCP StructuredTools), we must run the
        # graph via `ainvoke` so ToolNode dispatches tools asynchronously.
        if self._has_async_only_tools():
            try:
                asyncio.get_running_loop()
            except RuntimeError:
                return asyncio.run(self._ainvoke(input, **config))

            raise RuntimeError(
                "This agent has async-only tools, but `.invoke()` was called "
                "from an async context (a running event loop was detected). "
                "Use `await agent.ainvoke(...)` instead."
            )

        try:
            return self.compiled_graph.invoke(
                input, config=config, context=self.context
            )
        except Exception as e:
            # Fallback: if a tool raises the canonical sync-invoke error, retry
            # with ainvoke for backwards compatibility.
            if "does not support sync invocation" in str(e):
                try:
                    asyncio.get_running_loop()
                except RuntimeError:
                    return asyncio.run(self._ainvoke(input, **config))
            raise

    async def _ainvoke(self, input, **config):
        config = self.build_config(**config)
        graph = await self._aget_async_compiled_graph()
        try:
            return await graph.ainvoke(
                input, config=config, context=self.context
            )
        finally:
            await self.aclose()

    def _stream(self, input, **config):
        config = self.build_config(**config)
        yield from self.compiled_graph.stream(
            input, config=config, context=self.context
        )

    def __call__(self, inputs: InputLike, /, **kwargs: Any) -> Any:
        """Specify calling behavior for class instance."""
        return self.invoke(inputs, **kwargs)

    # Runtime enforcement: forbid subclasses from overriding invoke
    def __init_subclass__(cls, **kwargs):
        """Ensure subclass does not override key method."""
        super().__init_subclass__(**kwargs)
        if "invoke" in cls.__dict__:
            err_msg = (
                f"{cls.__name__} must not override BaseAgent.invoke(); "
                "implement _invoke() only."
            )
            raise TypeError(err_msg)

        # Init graph after subclass has been fully constructed
        orig_init = cls.__init__

        def __init__(self, *args, **kwargs):
            orig_init(self, *args, **kwargs)
            self.__post_init__()

        cls.__init__ = __init__

    def __post_init__(self):
        self.build_graph()

    def stream(
        self,
        inputs: InputLike,
        config: Any | None = None,  # allow positional/keyword like LangGraph
        /,
        *,
        raw_debug: bool = False,
        save_json: bool | None = None,
        save_otel: bool | None = None,
        metrics_path: str | None = None,
        save_raw_snapshot: bool | None = None,
        save_raw_records: bool | None = None,
        **kwargs: Any,
    ) -> Iterator[Any]:
        """Streams agent responses with telemetry tracking.

        This method serves as the public streaming entry point for agent interactions.
        It wraps the actual streaming implementation with telemetry tracking to capture
        metrics and debugging information.

        Args:
            inputs: The input to process, which will be normalized internally.
            config: Optional configuration for the agent, compatible with LangGraph
                positional/keyword argument style.
            raw_debug: If True, renders raw debug information in telemetry output.
            save_json: If True, saves telemetry data as JSON.
            metrics_path: Optional file path where metrics should be saved.
            save_raw_snapshot: If True, saves raw snapshot data in telemetry.
            save_raw_records: If True, saves raw record data in telemetry.
            **kwargs: Additional keyword arguments passed to the streaming
                implementation.

        Returns:
            An iterator yielding the agent's responses.

        Note:
            This method tracks invocation depth to properly handle nested agent calls
            and ensure telemetry is only rendered once at the top level.
        """
        # Track invocation depth to handle nested agent calls
        BaseAgent._invoke_depth += 1

        try:
            # Start telemetry tracking for top-level invocations only
            if BaseAgent._invoke_depth == 1:
                self.telemetry.begin_run(
                    agent=self.name, thread_id=self.thread_id
                )

            # Normalize inputs and delegate to the actual streaming implementation
            normalized = self._normalize_inputs(inputs)
            stream_config = dict(config or {})
            yield from self._stream(normalized, **stream_config, **kwargs)

        finally:
            # Decrement invocation depth when exiting
            BaseAgent._invoke_depth -= 1

            # Render telemetry data only for top-level invocations
            if BaseAgent._invoke_depth == 0:
                self.telemetry.render(
                    raw=raw_debug,
                    save_json=save_json,
                    save_otel=save_otel,
                    filepath=metrics_path,
                    save_raw_snapshot=save_raw_snapshot,
                    save_raw_records=save_raw_records,
                )

    def _default_node_tags(
        self, name: str, extra: Sequence[str] | None = None
    ) -> list[str]:
        """Generate default tags for a graph node.

        Args:
            name: The name of the node.
            extra: Optional sequence of additional tags to include.

        Returns:
            list[str]: A list of tags for the node, including the agent name, 'graph',
                the node name, and any extra tags provided.
        """
        # Start with standard tags: agent name, graph indicator, and node name
        tags = [self.name, "graph", name]

        # Add any extra tags if provided
        if extra:
            tags.extend(extra)

        return tags

    def _node_cfg(self, name: str, *extra_tags: str) -> dict:
        """Build a consistent configuration for a node/runnable.

        Creates a configuration dict that can be reapplied after operations like
        .map(), subgraph compile, etc.

        Args:
            name: The name of the node.
            *extra_tags: Additional tags to include in the node configuration.

        Returns:
            dict: A configuration dictionary with run_name, tags, and metadata.
        """
        # Determine the namespace - use first extra tag if available, otherwise
        # convert agent name to snake_case
        ns = extra_tags[0] if extra_tags else _to_snake(self.name)

        # Combine all tags: agent name, graph indicator, node name, and any extra tags
        tags = [self.name, "graph", name, *extra_tags]

        # Return the complete configuration dictionary
        return dict(
            run_name="node",  # keep "node:" prefixing in the timer
            tags=tags,
            metadata={
                "langgraph_node": name,
                "ursa_ns": ns,
                "ursa_agent": self.name,
            },
        )

    def ns(self, runnable_or_fn, name: str, *extra_tags: str):
        """Return a runnable with node configuration applied.

        Applies the agent's node configuration to a runnable or callable. This method
        should be called again after operations like .map() or subgraph .compile() as
        these operations may drop configuration.

        Args:
            runnable_or_fn: A runnable or callable to configure.
            name: The name to assign to this node.
            *extra_tags: Additional tags to apply to the node.

        Returns:
            A configured runnable with the agent's node configuration applied.
        """
        # Convert input to a runnable if it's not already one
        r = coerce_to_runnable(runnable_or_fn, name=name, trace=True)
        # Apply node configuration and return the configured runnable
        return r.with_config(**self._node_cfg(name, *extra_tags))

    def _wrap_node(self, fn_or_runnable, name: str, *extra_tags: str):
        """Wrap a function or runnable as a node in the graph.

        This is a convenience wrapper around the ns() method.

        Args:
            fn_or_runnable: A function or runnable to wrap as a node.
            name: The name to assign to this node.
            *extra_tags: Additional tags to apply to the node.

        Returns:
            A configured runnable with the agent's node configuration applied.
        """
        return self.ns(fn_or_runnable, name, *extra_tags)

    def _wrap_cond(self, fn: Any, name: str, *extra_tags: str):
        """Wrap a conditional function as a routing node in the graph.

        Creates a runnable lambda with routing-specific configuration.

        Args:
            fn: The conditional function to wrap.
            name: The name of the routing node.
            *extra_tags: Additional tags to apply to the node.

        Returns:
            A configured RunnableLambda with routing-specific metadata.
        """
        # Use the first extra tag as namespace, or fall back to agent name in snake_case
        ns = extra_tags[0] if extra_tags else _to_snake(self.name)

        # Create and return a configured RunnableLambda for routing
        return RunnableLambda(fn).with_config(
            run_name="node",
            tags=[
                self.name,
                "graph",
                f"route:{name}",
                *extra_tags,
            ],
            metadata={
                "langgraph_node": f"route:{name}",
                "ursa_ns": ns,
                "ursa_agent": self.name,
            },
        )

    def _named(self, runnable: Any, name: str, *extra_tags: str):
        """Apply a specific name and configuration to a runnable.

        Configures a runnable with a specific name and the agent's metadata.

        Args:
            runnable: The runnable to configure.
            name: The name to assign to this runnable.
            *extra_tags: Additional tags to apply to the runnable.

        Returns:
            A configured runnable with the specified name and agent metadata.
        """
        # Use the first extra tag as namespace, or fall back to agent name in snake_case
        ns = extra_tags[0] if extra_tags else _to_snake(self.name)

        # Apply configuration and return the configured runnable
        return runnable.with_config(
            run_name=name,
            tags=[self.name, "graph", name, *extra_tags],
            metadata={
                "langgraph_node": name,
                "ursa_ns": ns,
                "ursa_agent": self.name,
            },
        )

compiled_graph cached property

Return the sync compiled StateGraph application for the agent.

context property

Immutable run-scoped information provided to the Agent's graph

name property

Agent name.

storage cached property

Create a sync SQLite-backed LangGraph store.

__call__(inputs, /, **kwargs)

Specify calling behavior for class instance.

Source code in src/ursa/agents/base.py

def __call__(self, inputs: InputLike, /, **kwargs: Any) -> Any:
    """Specify calling behavior for class instance."""
    return self.invoke(inputs, **kwargs)

__init__(llm, workspace=None, agent_name=None, group='default', checkpointer=None, enable_metrics=True, metrics_dir='ursa_metrics', autosave_metrics=True, otel_metrics=False, thread_id=None, tokens_before_summarize=50000, messages_to_keep=20, max_single_tool_message_tokens=100000, rag_tools=None, rag_tool_group=None, rag_tool_embedding=None, rag_tool_return_k=10)

Initializes the base agent with a language model and optional configurations.

Parameters:

Name	Type	Description	Default
llm	BaseChatModel	a BaseChatModel instance.	required
checkpointer	Optional[BaseCheckpointSaver]	Optional checkpoint saver for persisting agent state.	None
enable_metrics	bool	Whether to collect performance and usage metrics.	True
metrics_dir	str	Directory path where metrics will be saved.	'ursa_metrics'
autosave_metrics	bool	Whether to automatically save metrics to disk.	True
thread_id	Optional[str]	Unique identifier for this agent instance. Generated if not provided.	None

Source code in src/ursa/agents/base.py

def __init__(
    self,
    llm: BaseChatModel,
    workspace: Optional[Path] = None,
    agent_name: Optional[str] = None,
    group: Optional[str] = "default",
    checkpointer: Optional[BaseCheckpointSaver] = None,
    enable_metrics: bool = True,
    metrics_dir: str = "ursa_metrics",  # dir to save metrics, with a default
    autosave_metrics: bool = True,
    otel_metrics: bool = False,
    thread_id: Optional[str] = None,
    tokens_before_summarize: int = 50000,
    messages_to_keep: int = 20,
    max_single_tool_message_tokens: int = 100000,
    rag_tools: str | Sequence[str] | None = None,
    rag_tool_group: str | None = None,
    rag_tool_embedding: Embeddings | None = None,
    rag_tool_return_k: int = 10,
):
    """Initializes the base agent with a language model and optional configurations.

    Args:
        llm: a BaseChatModel instance.
        checkpointer: Optional checkpoint saver for persisting agent state.
        enable_metrics: Whether to collect performance and usage metrics.
        metrics_dir: Directory path where metrics will be saved.
        autosave_metrics: Whether to automatically save metrics to disk.
        thread_id: Unique identifier for this agent instance. Generated if not
                   provided.
    """
    self.llm: BaseChatModel = llm
    self.workspace = Path(workspace or ".")
    self.agent_name = agent_name
    self.group = validate_group_name(group)
    self.tokens_before_summarize = tokens_before_summarize
    self.messages_to_keep = messages_to_keep
    self.max_single_tool_message_tokens = max_single_tool_message_tokens
    from ursa.rag.persistence import normalize_rag_tool_names

    self.rag_tools = tuple(normalize_rag_tool_names(rag_tools))
    self.rag_tool_group = rag_tool_group
    self.rag_tool_embedding = rag_tool_embedding
    self.rag_tool_return_k = rag_tool_return_k
    enforce_model_group_policy(self.llm, self.group)
    if (
        not group_root_dir(self.group).exists()
        and self.group != DEFAULT_GROUP_NAME
    ):
        raise ValueError(
            (
                f"Group '{self.group}' does not exist. "
                f"Please use `ursa create-group {self.group} <group_config_file>` to create"
            )
        )
    set_checkpointer = True if checkpointer else False
    set_name = True if agent_name else False
    self.checkpointer = checkpointer
    self._checkpointer_was_provided = set_checkpointer
    self._async_checkpointer: BaseCheckpointSaver | None = None
    self._async_storage: BaseStore | None = None
    self._async_compiled_graph: CompiledStateGraph | None = None
    persist_agent = (agent_name is not None) or set_checkpointer
    if persist_agent:
        if set_checkpointer:
            if set_name:
                print(
                    (
                        "[WARNING]: Both checkpointer and den persistence set."
                        " Using checkpointer, but only use one in the future."
                    )
                )
            self.den = self.workspace
        else:
            den_name = Path(self.group) / "agents" / agent_name
            self.den = group_agents_dir(self.group) / agent_name
            if not self.den.exists():
                print(f"[Agent Created]: {den_name}")
            self.checkpointer = Checkpointer.from_workspace(self.den)
    else:
        # Keep current behavior if the user is not persisting.
        self.den = self.workspace
    self.thread_id = thread_id or "ursa"
    self.telemetry = Telemetry(
        enable=enable_metrics,
        output_dir=self.den.joinpath(metrics_dir),
        save_json_default=autosave_metrics,
    )

    self.workspace.mkdir(exist_ok=True, parents=True)
    self.den.mkdir(exist_ok=True, parents=True)

__init_subclass__(**kwargs)

Ensure subclass does not override key method.

Source code in src/ursa/agents/base.py

def __init_subclass__(cls, **kwargs):
    """Ensure subclass does not override key method."""
    super().__init_subclass__(**kwargs)
    if "invoke" in cls.__dict__:
        err_msg = (
            f"{cls.__name__} must not override BaseAgent.invoke(); "
            "implement _invoke() only."
        )
        raise TypeError(err_msg)

    # Init graph after subclass has been fully constructed
    orig_init = cls.__init__

    def __init__(self, *args, **kwargs):
        orig_init(self, *args, **kwargs)
        self.__post_init__()

    cls.__init__ = __init__

aclose() async

Close async SQLite resources owned by this agent when possible.

Source code in src/ursa/agents/base.py

async def aclose(self) -> None:
    """Close async SQLite resources owned by this agent when possible."""
    if self._async_storage is not None:
        await self._async_storage.__aexit__(None, None, None)
        await self._async_storage.conn.close()
        self._async_storage = None
        self._async_compiled_graph = None
    if isinstance(self._async_checkpointer, AsyncSqliteSaver):
        await self._async_checkpointer.conn.close()
        self._async_checkpointer = None
        self._async_compiled_graph = None

add_node(f, node_name=None, agent_name=None, **kwargs)

Add a node to the state graph with token usage tracking.

This method adds a function as a node to the state graph, wrapping it to track token usage during execution. The node is identified by either the provided node_name or the function's name.

Parameters:

Name	Type	Description	Default
f	Callable[..., Mapping[str, Any]]	The function to add as a node. Should return a mapping of string keys to any values.	required
node_name	Optional[str]	Optional name for the node. If not provided, the function's name will be used.	None
agent_name	Optional[str]	Optional agent name for tracking. If not provided, the agent's name in snake_case will be used.	None

Returns:

Type	Description
StateGraph	The updated StateGraph with the new node added.

Source code in src/ursa/agents/base.py

def add_node(
    self,
    f: Callable[..., Mapping[str, Any]],
    node_name: Optional[str] = None,
    agent_name: Optional[str] = None,
    **kwargs,
) -> StateGraph:
    """Add a node to the state graph with token usage tracking.

    This method adds a function as a node to the state graph, wrapping it to track
    token usage during execution. The node is identified by either the provided
    node_name or the function's name.

    Args:
        f: The function to add as a node. Should return a mapping of string keys to
            any values.
        node_name: Optional name for the node. If not provided, the function's name
            will be used.
        agent_name: Optional agent name for tracking. If not provided, the agent's
            name in snake_case will be used.

    Returns:
        The updated StateGraph with the new node added.
    """
    _node_name = node_name or f.__name__
    _agent_name = agent_name or _to_snake(self.name)
    wrapped_node = self._wrap_node(f, _node_name, _agent_name)
    return self.graph.add_node(_node_name, wrapped_node, **kwargs)

ainvoke(inputs=None, /, *, raw_debug=False, save_json=None, save_otel=None, metrics_path=None, save_raw_snapshot=None, save_raw_records=None, config=None, **kwargs) async

Asynchrnously executes the agent with the provided inputs and configuration.

(Async version of invoke.)

This is the main entry point for agent execution. It handles input normalization, telemetry tracking, and proper execution context management. The method supports flexible input formats - either as a positional argument or as keyword arguments.

Parameters:

Name	Type	Description	Default
inputs	Optional[InputLike]	Optional positional input to the agent. If provided, all non-control keyword arguments will be rejected to avoid ambiguity.	None
raw_debug	bool	If True, displays raw telemetry data for debugging purposes.	False
save_json	Optional[bool]	If True, saves telemetry data as JSON.	None
save_otel	Optional[bool]	If True, saves telemetry data to OpenTelemetry endpoint.	None
metrics_path	Optional[str]	Optional file path where telemetry metrics should be saved.	None
save_raw_snapshot	Optional[bool]	If True, saves a raw snapshot of the telemetry data.	None
save_raw_records	Optional[bool]	If True, saves raw telemetry records.	None
config	Optional[dict]	Optional configuration dictionary to override default settings.	None
**kwargs	Any	Additional keyword arguments that can be either: - Input parameters (when no positional input is provided) - Control parameters recognized by the agent	{}

Returns:

Type	Description
Any	The result of the agent's execution.

Raises:

Type	Description
TypeError	If both positional inputs and non-control keyword arguments are provided simultaneously.

Source code in src/ursa/agents/base.py

@final
async def ainvoke(
    self,
    inputs: Optional[InputLike] = None,
    /,
    *,
    raw_debug: bool = False,
    save_json: Optional[bool] = None,
    save_otel: Optional[bool] = None,
    metrics_path: Optional[str] = None,
    save_raw_snapshot: Optional[bool] = None,
    save_raw_records: Optional[bool] = None,
    config: Optional[dict] = None,
    **kwargs: Any,
) -> Any:
    """Asynchrnously executes the agent with the provided inputs and configuration.

    (Async version of `invoke`.)

    This is the main entry point for agent execution. It handles input normalization,
    telemetry tracking, and proper execution context management. The method supports
    flexible input formats - either as a positional argument or as keyword arguments.

    Args:
        inputs: Optional positional input to the agent. If provided, all non-control
            keyword arguments will be rejected to avoid ambiguity.
        raw_debug: If True, displays raw telemetry data for debugging purposes.
        save_json: If True, saves telemetry data as JSON.
        save_otel: If True, saves telemetry data to OpenTelemetry endpoint.
        metrics_path: Optional file path where telemetry metrics should be saved.
        save_raw_snapshot: If True, saves a raw snapshot of the telemetry data.
        save_raw_records: If True, saves raw telemetry records.
        config: Optional configuration dictionary to override default settings.
        **kwargs: Additional keyword arguments that can be either:
            - Input parameters (when no positional input is provided)
            - Control parameters recognized by the agent

    Returns:
        The result of the agent's execution.

    Raises:
        TypeError: If both positional inputs and non-control keyword arguments are
            provided simultaneously.
    """
    return await self._ainvoke_engine(
        invoke_method=self._ainvoke,
        inputs=inputs,
        raw_debug=raw_debug,
        save_json=save_json,
        save_otel=save_otel,
        metrics_path=metrics_path,
        save_raw_snapshot=save_raw_snapshot,
        save_raw_records=save_raw_records,
        config=config,
        **kwargs,
    )

build_config(**overrides)

Constructs a config dictionary for agent operations with telemetry support.

This method creates a standardized configuration dictionary that includes thread identification, telemetry callbacks, and other metadata needed for agent operations. The configuration can be customized through override parameters.

Parameters:

Name	Type	Description	Default
**overrides		Optional configuration overrides that can include keys like 'recursion_limit', 'configurable', 'metadata', 'tags', etc.	{}

Returns:

Name	Type	Description
dict	dict	A complete configuration dictionary with all necessary parameters.

Source code in src/ursa/agents/base.py

def build_config(self, **overrides) -> dict:
    """Constructs a config dictionary for agent operations with telemetry support.

    This method creates a standardized configuration dictionary that includes thread
    identification, telemetry callbacks, and other metadata needed for agent
    operations. The configuration can be customized through override parameters.

    Args:
        **overrides: Optional configuration overrides that can include keys like
            'recursion_limit', 'configurable', 'metadata', 'tags', etc.

    Returns:
        dict: A complete configuration dictionary with all necessary parameters.
    """
    # Create the base configuration with essential fields.
    base = {
        "configurable": {"thread_id": self.thread_id},
        "metadata": {
            "thread_id": self.thread_id,
            "telemetry_run_id": self.telemetry.context.get("run_id"),
        },
        "tags": [self.name],
        "callbacks": [
            *self.telemetry.callbacks,
            DEFAULT_EVENT_LOGGING_HANDLER,
        ],
    }

    # Try to determine the model name from either direct or nested attributes
    model_name = getattr(self, "llm_model", None) or getattr(
        getattr(self, "llm", None), "model", None
    )

    # Add model name to metadata if available
    if model_name:
        base["metadata"]["model"] = model_name

    # Handle configurable dictionary overrides by merging with base configurable
    if "configurable" in overrides and isinstance(
        overrides["configurable"], dict
    ):
        base["configurable"].update(overrides.pop("configurable"))

    # Handle metadata dictionary overrides by merging with base metadata
    if "metadata" in overrides and isinstance(overrides["metadata"], dict):
        base["metadata"].update(overrides.pop("metadata"))

    # Merge tags from caller-provided overrides, avoid duplicates
    if "tags" in overrides and isinstance(overrides["tags"], list):
        base["tags"] = base["tags"] + [
            t for t in overrides.pop("tags") if t not in base["tags"]
        ]

    if "callbacks" in overrides:
        callbacks = merge_configs(
            {"callbacks": base["callbacks"]},
            {"callbacks": overrides.pop("callbacks")},
        ).get("callbacks")
        if isinstance(callbacks, list):
            callbacks = self._dedupe_callbacks(callbacks)
        base["callbacks"] = callbacks

    # Apply any remaining overrides directly to the base configuration
    base.update(overrides)

    return base

build_graph()

Build and return the StateGraph backing this agent.

Source code in src/ursa/agents/base.py

@final
def build_graph(self) -> StateGraph:
    """Build and return the StateGraph backing this agent."""
    self.graph = StateGraph(
        self.state_type,
        context_schema=AgentContext,
    )
    self._build_graph()
    return self.graph

close()

Close sync SQLite resources owned by this agent when possible.

Source code in src/ursa/agents/base.py

def close(self) -> None:
    """Close sync SQLite resources owned by this agent when possible."""
    if "storage" in self.__dict__:
        self.storage.conn.close()
    if isinstance(self.checkpointer, SqliteSaver):
        self.checkpointer.conn.close()

events(config=None)

Return a helper for emitting structured agent progress events.

Source code in src/ursa/agents/base.py

def events(self, config: RunnableConfig | None = None) -> AgentEvents:
    """Return a helper for emitting structured agent progress events."""
    return AgentEvents(agent=self.name, config=config)

format_query(prompt, state=None)

Format a plain text prompt into the agent's input schema possibly incorporating the prior state.

Agents should override this method for their operation

Source code in src/ursa/agents/base.py

def format_query(self, prompt: str, state: TState | None = None) -> TState:
    """Format a plain text prompt into the agent's input schema
    possibly incorporating the prior state.

    Agents should override this method for their operation
    """

    if state is not None and "messages" in state:
        state["messages"].append(HumanMessage(content=str(prompt)))
        return state
    return self._normalize_inputs(prompt)

format_result(result)

Extracts a plain text response from the Agent's output schema

Agents should override this method for their operation

Source code in src/ursa/agents/base.py

def format_result(self, result: TState) -> str:
    """Extracts a plain text response from the Agent's output schema

    Agents should override this method for their operation
    """

    if "messages" in result:
        if isinstance(result["messages"], list) and isinstance(
            result["messages"][-1], BaseMessage
        ):
            return result["messages"][-1].text
    raise NotImplementedError()

invoke(inputs=None, /, *, raw_debug=False, save_json=None, save_otel=None, metrics_path=None, save_raw_snapshot=None, save_raw_records=None, config=None, **kwargs)

Executes the agent with the provided inputs and configuration.

This is the main entry point for agent execution. It handles input normalization, telemetry tracking, and proper execution context management. The method supports flexible input formats - either as a positional argument or as keyword arguments.

Parameters:

Name	Type	Description	Default
inputs	Optional[InputLike]	Optional positional input to the agent. If provided, all non-control keyword arguments will be rejected to avoid ambiguity.	None
raw_debug	bool	If True, displays raw telemetry data for debugging purposes.	False
save_json	Optional[bool]	If True, saves telemetry data as JSON.	None
save_otel	Optional[bool]	If True, saves telemetry data to OpenTelemetry endpoint.	None
metrics_path	Optional[str]	Optional file path where telemetry metrics should be saved.	None
save_raw_snapshot	Optional[bool]	If True, saves a raw snapshot of the telemetry data.	None
save_raw_records	Optional[bool]	If True, saves raw telemetry records.	None
config	Optional[dict]	Optional configuration dictionary to override default settings.	None
**kwargs	Any	Additional keyword arguments that can be either: - Input parameters (when no positional input is provided) - Control parameters recognized by the agent	{}

Returns:

Type	Description
Any	The result of the agent's execution.

Raises:

Type	Description
TypeError	If both positional inputs and non-control keyword arguments are provided simultaneously.

Source code in src/ursa/agents/base.py

@final
def invoke(
    self,
    inputs: Optional[InputLike] = None,
    /,
    *,
    raw_debug: bool = False,
    save_json: Optional[bool] = None,
    save_otel: Optional[bool] = None,
    metrics_path: Optional[str] = None,
    save_raw_snapshot: Optional[bool] = None,
    save_raw_records: Optional[bool] = None,
    config: Optional[dict] = None,
    **kwargs: Any,
) -> Any:
    """Executes the agent with the provided inputs and configuration.

    This is the main entry point for agent execution. It handles input normalization,
    telemetry tracking, and proper execution context management. The method supports
    flexible input formats - either as a positional argument or as keyword arguments.

    Args:
        inputs: Optional positional input to the agent. If provided, all non-control
            keyword arguments will be rejected to avoid ambiguity.
        raw_debug: If True, displays raw telemetry data for debugging purposes.
        save_json: If True, saves telemetry data as JSON.
        save_otel: If True, saves telemetry data to OpenTelemetry endpoint.
        metrics_path: Optional file path where telemetry metrics should be saved.
        save_raw_snapshot: If True, saves a raw snapshot of the telemetry data.
        save_raw_records: If True, saves raw telemetry records.
        config: Optional configuration dictionary to override default settings.
        **kwargs: Additional keyword arguments that can be either:
            - Input parameters (when no positional input is provided)
            - Control parameters recognized by the agent

    Returns:
        The result of the agent's execution.

    Raises:
        TypeError: If both positional inputs and non-control keyword arguments are
            provided simultaneously.
    """
    return self._invoke_engine(
        invoke_method=self._invoke,
        inputs=inputs,
        raw_debug=raw_debug,
        save_json=save_json,
        save_otel=save_otel,
        metrics_path=metrics_path,
        save_raw_snapshot=save_raw_snapshot,
        save_raw_records=save_raw_records,
        config=config,
        **kwargs,
    )

messages_update(state, message_update, *, full_overwrite=False, extra=None)

Build a LangGraph-safe update for states with a messages reducer.

Source code in src/ursa/agents/base.py

def messages_update(
    self,
    state: Mapping[str, Any],
    message_update: Any,
    *,
    full_overwrite: bool = False,
    extra: Mapping[str, Any] | None = None,
) -> dict[str, Any]:
    """Build a LangGraph-safe update for states with a ``messages`` reducer."""
    update = dict(extra or {})
    if full_overwrite:
        messages = list(state.get("messages") or [])
        if message_update is None:
            update["messages"] = Overwrite(messages)
        else:
            if isinstance(message_update, list):
                messages.extend(message_update)
            else:
                messages.append(message_update)
            update["messages"] = Overwrite(messages)
    else:
        update["messages"] = message_update
    return update

ns(runnable_or_fn, name, *extra_tags)

Return a runnable with node configuration applied.

Applies the agent's node configuration to a runnable or callable. This method should be called again after operations like .map() or subgraph .compile() as these operations may drop configuration.

Parameters:

Name	Type	Description	Default
runnable_or_fn		A runnable or callable to configure.	required
name	str	The name to assign to this node.	required
*extra_tags	str	Additional tags to apply to the node.	()

Returns:

Type	Description
	A configured runnable with the agent's node configuration applied.

Source code in src/ursa/agents/base.py

def ns(self, runnable_or_fn, name: str, *extra_tags: str):
    """Return a runnable with node configuration applied.

    Applies the agent's node configuration to a runnable or callable. This method
    should be called again after operations like .map() or subgraph .compile() as
    these operations may drop configuration.

    Args:
        runnable_or_fn: A runnable or callable to configure.
        name: The name to assign to this node.
        *extra_tags: Additional tags to apply to the node.

    Returns:
        A configured runnable with the agent's node configuration applied.
    """
    # Convert input to a runnable if it's not already one
    r = coerce_to_runnable(runnable_or_fn, name=name, trace=True)
    # Apply node configuration and return the configured runnable
    return r.with_config(**self._node_cfg(name, *extra_tags))

prepare_messages_context(state, *, system_prompt=None, summary_prompt=None, patch_dangling=True, summarize=True)

Apply standard message-history maintenance before an LLM call.

Returns (new_state, full_overwrite_required). Graph nodes using the add_messages reducer should use :meth:messages_update with the returned flag to avoid appending duplicate history after summarization or dangling-tool patching.

Source code in src/ursa/agents/base.py

def prepare_messages_context(
    self,
    state: Mapping[str, Any],
    *,
    system_prompt: str | None = None,
    summary_prompt: str | None = None,
    patch_dangling: bool = True,
    summarize: bool = True,
) -> tuple[Mapping[str, Any], bool]:
    """Apply standard message-history maintenance before an LLM call.

    Returns ``(new_state, full_overwrite_required)``. Graph nodes using the
    ``add_messages`` reducer should use :meth:`messages_update` with the
    returned flag to avoid appending duplicate history after summarization or
    dangling-tool patching.
    """
    new_state = deepcopy(state)
    full_overwrite = False
    if summarize:
        new_state, full_overwrite = self._summarize_context(
            new_state,
            system_prompt=system_prompt,
            summary_prompt=summary_prompt,
        )
    if patch_dangling:
        new_state, full_overwrite = self._patch_dangling(
            new_state, full_overwrite
        )
    return new_state, full_overwrite

stream(inputs, config=None, /, *, raw_debug=False, save_json=None, save_otel=None, metrics_path=None, save_raw_snapshot=None, save_raw_records=None, **kwargs)

Streams agent responses with telemetry tracking.

This method serves as the public streaming entry point for agent interactions. It wraps the actual streaming implementation with telemetry tracking to capture metrics and debugging information.

Parameters:

Name	Type	Description	Default
inputs	InputLike	The input to process, which will be normalized internally.	required
config	Any | None	Optional configuration for the agent, compatible with LangGraph positional/keyword argument style.	None
raw_debug	bool	If True, renders raw debug information in telemetry output.	False
save_json	bool | None	If True, saves telemetry data as JSON.	None
metrics_path	str | None	Optional file path where metrics should be saved.	None
save_raw_snapshot	bool | None	If True, saves raw snapshot data in telemetry.	None
save_raw_records	bool | None	If True, saves raw record data in telemetry.	None
**kwargs	Any	Additional keyword arguments passed to the streaming implementation.	{}

Returns:

Type	Description
Iterator[Any]	An iterator yielding the agent's responses.

Note

This method tracks invocation depth to properly handle nested agent calls and ensure telemetry is only rendered once at the top level.

Source code in src/ursa/agents/base.py

def stream(
    self,
    inputs: InputLike,
    config: Any | None = None,  # allow positional/keyword like LangGraph
    /,
    *,
    raw_debug: bool = False,
    save_json: bool | None = None,
    save_otel: bool | None = None,
    metrics_path: str | None = None,
    save_raw_snapshot: bool | None = None,
    save_raw_records: bool | None = None,
    **kwargs: Any,
) -> Iterator[Any]:
    """Streams agent responses with telemetry tracking.

    This method serves as the public streaming entry point for agent interactions.
    It wraps the actual streaming implementation with telemetry tracking to capture
    metrics and debugging information.

    Args:
        inputs: The input to process, which will be normalized internally.
        config: Optional configuration for the agent, compatible with LangGraph
            positional/keyword argument style.
        raw_debug: If True, renders raw debug information in telemetry output.
        save_json: If True, saves telemetry data as JSON.
        metrics_path: Optional file path where metrics should be saved.
        save_raw_snapshot: If True, saves raw snapshot data in telemetry.
        save_raw_records: If True, saves raw record data in telemetry.
        **kwargs: Additional keyword arguments passed to the streaming
            implementation.

    Returns:
        An iterator yielding the agent's responses.

    Note:
        This method tracks invocation depth to properly handle nested agent calls
        and ensure telemetry is only rendered once at the top level.
    """
    # Track invocation depth to handle nested agent calls
    BaseAgent._invoke_depth += 1

    try:
        # Start telemetry tracking for top-level invocations only
        if BaseAgent._invoke_depth == 1:
            self.telemetry.begin_run(
                agent=self.name, thread_id=self.thread_id
            )

        # Normalize inputs and delegate to the actual streaming implementation
        normalized = self._normalize_inputs(inputs)
        stream_config = dict(config or {})
        yield from self._stream(normalized, **stream_config, **kwargs)

    finally:
        # Decrement invocation depth when exiting
        BaseAgent._invoke_depth -= 1

        # Render telemetry data only for top-level invocations
        if BaseAgent._invoke_depth == 0:
            self.telemetry.render(
                raw=raw_debug,
                save_json=save_json,
                save_otel=save_otel,
                filepath=metrics_path,
                save_raw_snapshot=save_raw_snapshot,
                save_raw_records=save_raw_records,
            )

write_state(filename, state)

Writes agent state to a JSON file.

Serializes the provided state dictionary to JSON format and writes it to the specified file. The JSON is written with non-ASCII characters preserved.

Parameters:

Name	Type	Description	Default
filename	str	Path to the file where state will be written.	required
state	dict	Dictionary containing the agent state to be serialized.	required

Source code in src/ursa/agents/base.py

def write_state(self, filename: str, state: dict) -> None:
    """Writes agent state to a JSON file.

    Serializes the provided state dictionary to JSON format and writes it to the
    specified file. The JSON is written with non-ASCII characters preserved.

    Args:
        filename: Path to the file where state will be written.
        state: Dictionary containing the agent state to be serialized.
    """
    json_state = dumps(state, ensure_ascii=False)
    with open(filename, "w") as f:
        f.write(json_state)

chat_agent

BasicChatAgent

Bases: BaseAgent[ChatState]

Basic Chat Agent

Source code in src/ursa/agents/chat_agent.py

class BasicChatAgent(BaseAgent[ChatState]):
    """Basic Chat Agent"""

    state_type = ChatState

    def _response_node(
        self, state: ChatState, runtime: Runtime[AgentContext]
    ) -> ChatState:
        new_state, full_overwrite = self.prepare_messages_context(state)
        res = self.llm.invoke(new_state["messages"])
        return self.messages_update(
            new_state, [res], full_overwrite=full_overwrite
        )

    def format_query(self, prompt: str, state: ChatState | None = None):
        if state is None:
            state = ChatState(
                messages=[SystemMessage(content=get_chatter_system_prompt())]
            )
        state["messages"].append(HumanMessage(content=prompt))

        return state

    def format_result(self, result: ChatState) -> str:
        return result["messages"][-1].text

    def _build_graph(self):
        self.add_node(self._response_node)
        self.graph.set_entry_point("_response_node")
        self.graph.set_finish_point("_response_node")

ChatAgent

Bases: AgentWithTools, BasicChatAgent

Chat Agent

Source code in src/ursa/agents/chat_agent.py

class ChatAgent(AgentWithTools, BasicChatAgent):
    """Chat Agent"""

    state_type = ChatState

    def __init__(
        self,
        llm: BaseChatModel,
        use_web: bool = False,
        **kwargs,
    ):
        default_tools = [
            run_command,
            write_code,
            edit_code,
            read_file,
            download_file_tool,
            read_image_tool,
            list_experiences,
            write_experience,
            read_experience,
            edit_experience,
        ]
        if use_web:
            default_tools.extend([
                run_web_search,
                run_osti_search,
                run_arxiv_search,
            ])
        super().__init__(llm=llm, tools=default_tools, **kwargs)

    def _build_graph(self):
        # Bind tools to llm and context summarizer
        self.llm = self.llm.bind_tools(self.tools.values())

        self.add_node(self._response_node, "respond")
        self.add_node(self.tool_node, "tool_node")
        self.graph.set_entry_point("respond")
        self.graph.add_conditional_edges(
            "respond",
            self._wrap_cond(should_continue, "should_continue"),
            {"continue": "tool_node", "finish": END},
        )
        self.graph.add_edge("tool_node", "respond")

should_continue(state)

Return 'finish' if no tool calls in the last message, else 'continue'.

Parameters:

Name	Type	Description	Default
state	ChatState	The current execution state containing messages.	required

Returns:

Type	Description
Literal['finish', 'continue']	A literal "finish" if the last message has no tool calls,
Literal['finish', 'continue']	otherwise "continue".

Source code in src/ursa/agents/chat_agent.py

def should_continue(state: ChatState) -> Literal["finish", "continue"]:
    """Return 'finish' if no tool calls in the last message, else 'continue'.

    Args:
        state: The current execution state containing messages.

    Returns:
        A literal "finish" if the last message has no tool calls,
        otherwise "continue".
    """
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no tool call, then we finish
    if not last_message.tool_calls:
        return "finish"
    # Otherwise if there is, we continue
    else:
        return "continue"

deep_review_agent

DeepReviewAgent

Bases: AgentWithTools, BaseAgent[DeepReviewState]

Iterative adversarial review agent with explicit autonomous tool use.

The old Deep Review implementation performed hidden DuckDuckGo searches inside each phase. This implementation routes all external information access through configured LangChain tools:

• workspace tools are always available (list_workspace_files, read_file), and operate on the user-selected workspace rather than the agent den;
• web/search tools are only available when use_web=True;
• persisted RAG tools are merged automatically by AgentWithTools when rag_tools=... is supplied to the base-agent constructor;
• callers may add further tools via extra_tools or the normal AgentWithTools.add_tool / MCP mechanisms.

Source code in src/ursa/agents/deep_review_agent.py

class DeepReviewAgent(AgentWithTools, BaseAgent[DeepReviewState]):
    """Iterative adversarial review agent with explicit autonomous tool use.

    The old Deep Review implementation performed hidden DuckDuckGo searches inside
    each phase. This implementation routes all external information access through
    configured LangChain tools:

    - workspace tools are always available (`list_workspace_files`, `read_file`),
      and operate on the user-selected workspace rather than the agent den;
    - web/search tools are only available when `use_web=True`;
    - persisted RAG tools are merged automatically by `AgentWithTools` when
      `rag_tools=...` is supplied to the base-agent constructor;
    - callers may add further tools via `extra_tools` or the normal
      `AgentWithTools.add_tool` / MCP mechanisms.
    """

    state_type = DeepReviewState

    def __init__(
        self,
        llm: BaseChatModel,
        max_iterations: int = 2,
        use_web: bool = False,
        extra_tools: list[BaseTool] | None = None,
        **kwargs,
    ):
        default_tools: list[BaseTool] = [list_workspace_files, read_file]
        if use_web:
            default_tools.extend([
                run_web_search,
                run_osti_search,
                run_arxiv_search,
            ])
        if extra_tools:
            default_tools.extend(extra_tools)

        super().__init__(llm=llm, tools=default_tools, **kwargs)
        self.hypothesizer_prompt = hypothesizer_prompt
        self.critic_prompt = critic_prompt
        self.competitor_prompt = competitor_prompt
        self.strllm = self.llm | StrOutputParser()
        self.max_iterations = max_iterations
        self.use_web = use_web
        self.extra_tools = extra_tools or []

    def _normalize_inputs(self, inputs) -> DeepReviewState:
        if isinstance(inputs, str):
            return DeepReviewState(
                question=inputs,
                question_search_query=self._normalize_search_query(inputs),
                max_iterations=self.max_iterations,
                current_iteration=0,
                agent1_solution=[],
                agent2_critiques=[],
                agent3_perspectives=[],
                visited_sites=set(),
                messages=[],
                active_phase="",
            )

        state = dict(cast(Mapping[str, Any], inputs))
        question = str(state.get("question") or state.get("query") or "")
        state.setdefault("question", question)
        state.setdefault(
            "question_search_query", self._normalize_search_query(question)
        )
        state.setdefault("max_iterations", self.max_iterations)
        state.setdefault("current_iteration", 0)
        state.setdefault("agent1_solution", [])
        state.setdefault("agent2_critiques", [])
        state.setdefault("agent3_perspectives", [])
        state.setdefault("visited_sites", set())
        state.setdefault("messages", [])
        state.setdefault("active_phase", "")
        return cast(DeepReviewState, state)

    def format_result(self, result: DeepReviewState) -> str:
        return result.get("solution", "Deep review failed to return a solution")

    def _normalize_search_query(self, query: str | None) -> str:
        words = " ".join(str(query or "").replace("\n", " ").split())
        words = words.strip(" \"'")
        return " ".join(words.split()[:8])

    @staticmethod
    def _message_text(message: Any) -> str:
        text = getattr(message, "text", None)
        if isinstance(text, str):
            return text
        content = getattr(message, "content", message)
        if isinstance(content, str):
            return content
        return str(content)

    @staticmethod
    def _extract_urls(text: str) -> set[str]:
        return set(re.findall(r"https?://[^\s)\]}>,]+", text or ""))

    def _collect_visited_sites(self, state: DeepReviewState) -> set[str]:
        visited: set[str] = set(state.get("visited_sites") or set())
        for message in state.get("messages") or []:
            if message.__class__.__name__ == "ToolMessage":
                visited.update(self._extract_urls(self._message_text(message)))
        return visited

    def _base_phase_instructions(self, phase_label: str) -> str:
        tool_names = sorted(self.tools)
        web_instruction = (
            "Web/search tools are available because use_web=True. Use them only "
            "when they would materially improve the review."
            if self.use_web
            else "No web/search tools are available because use_web=False. Do not "
            "claim that you searched the web; rely on the prompt, workspace files, "
            "and any configured RAG tools."
        )
        return (
            f"You are the {phase_label} in an iterative deep-review process.\n"
            "You may autonomously call the configured tools when useful. Workspace "
            "documents live in the current workspace, not in the agent den. Use "
            "list_workspace_files to discover relevant files and read_file to read "
            "specific workspace documents. If RAG tools are available, use them for "
            "focused retrieval from configured document collections.\n"
            f"{web_instruction}\n"
            f"Available tools: {', '.join(tool_names) if tool_names else 'none'}.\n"
            "When you have enough information, answer directly without tool calls."
        )

    def _phase_prompt(self, phase: str, state: DeepReviewState) -> str:
        iteration = state["current_iteration"] + 1
        question = state["question"]
        if phase == "agent1":
            user_content = (
                f"Question: {question}\n"
                f"Deep-review iteration: {iteration}/{state['max_iterations']}\n"
            )
            if state.get("current_iteration", 0) > 0:
                user_content += (
                    f"\nPrevious solution: {state['agent1_solution'][-1]}"
                    f"\nCritique: {state['agent2_critiques'][-1]}"
                    "\nCompetitor/stakeholder perspective: "
                    f"{state['agent3_perspectives'][-1]}"
                    "\n\nExplicitly list how the new solution differs from the "
                    "previous solution, point by point, explaining what changes "
                    "were made in response to the critique and stakeholder "
                    "perspective. Then provide the updated solution."
                )
            else:
                user_content += (
                    "Generate an initial solution/hypothesis. Use available "
                    "workspace, RAG, or web tools only if they are relevant and "
                    "configured."
                )
            return user_content

        if phase == "agent2":
            solution = state["agent1_solution"][-1]
            return (
                f"Question: {question}\n"
                f"Proposed solution: {solution}\n"
                "Provide a detailed critique of this solution. Identify "
                "potential flaws, assumptions, missing evidence, and areas for "
                "improvement. Use available tools for targeted verification if "
                "useful and configured."
            )

        if phase == "agent3":
            solution = state["agent1_solution"][-1]
            critique = state["agent2_critiques"][-1]
            return (
                f"Question: {question}\n"
                f"Proposed solution: {solution}\n"
                f"Critique: {critique}\n"
                "Simulate how a competitor, government agency, stakeholder, or "
                "other adversarial party might respond to this solution. Use "
                "available tools for targeted context if useful and configured."
            )

        raise ValueError(f"Unknown deep-review phase: {phase}")

    def _role_system_prompt(self, phase: str) -> str:
        if phase == "agent1":
            role_prompt = self.hypothesizer_prompt
            label = "solution generator"
        elif phase == "agent2":
            role_prompt = self.critic_prompt
            label = "critic"
        elif phase == "agent3":
            role_prompt = self.competitor_prompt
            label = "competitor/stakeholder simulator"
        else:
            raise ValueError(f"Unknown deep-review phase: {phase}")
        return f"{role_prompt}\n\n{self._base_phase_instructions(label)}"

    def _role_output_key(self, phase: str) -> str:
        return {
            "agent1": "agent1_solution",
            "agent2": "agent2_critiques",
            "agent3": "agent3_perspectives",
        }[phase]

    def _run_role_phase(
        self,
        phase: str,
        state: DeepReviewState,
        config: RunnableConfig | None = None,
    ) -> DeepReviewState:
        events = self.events(config)
        iteration = state["current_iteration"] + 1
        events.emit(
            f"Deep-review {phase} pass {iteration}/{state['max_iterations']}",
            stage=phase,
            iteration=iteration,
            max_iterations=state["max_iterations"],
            tools=sorted(self.tools),
        )

        messages = list(state.get("messages") or [])
        new_phase_messages: list[BaseMessage] = []
        if state.get("active_phase") != phase:
            new_phase_messages = [
                SystemMessage(content=self._role_system_prompt(phase)),
                HumanMessage(content=self._phase_prompt(phase, state)),
            ]
            messages.extend(new_phase_messages)

        try:
            response = self.tool_llm.invoke(
                messages,
                self.build_config(tags=["deep_review", phase]),
            )
        except Exception as exc:  # noqa: BLE001
            response = AIMessage(content=f"Deep-review phase error: {exc}")
            events.emit(
                "Deep-review phase failed",
                stage=f"{phase}_error",
                error_type=type(exc).__name__,
                error=str(exc),
            )

        update: dict[str, Any] = {
            "messages": [*new_phase_messages, response],
            "active_phase": phase,
        }
        if getattr(response, "tool_calls", None):
            events.emit(
                "Deep-review phase requested tools",
                stage=f"{phase}_tool_request",
                tool_calls=[call.get("name") for call in response.tool_calls],
            )
            return cast(DeepReviewState, update)

        text = self._message_text(response)
        update[self._role_output_key(phase)] = [text]
        update["active_phase"] = ""
        update["visited_sites"] = self._collect_visited_sites(state)
        if phase == "agent1":
            update["question_search_query"] = self._normalize_search_query(
                state.get("question_search_query") or state.get("question")
            )
        events.emit(
            "Deep-review phase complete",
            stage=f"{phase}_result",
            preview=text,
        )
        return cast(DeepReviewState, update)

    def agent1_generate_solution(
        self,
        state: DeepReviewState,
        runtime: Runtime[AgentContext] | None = None,
        config: RunnableConfig | None = None,
    ) -> DeepReviewState:
        """Agent 1: autonomous solution generator."""
        return self._run_role_phase("agent1", state, config)

    def agent2_critique(
        self,
        state: DeepReviewState,
        runtime: Runtime[AgentContext] | None = None,
        config: RunnableConfig | None = None,
    ) -> DeepReviewState:
        """Agent 2: autonomous critic."""
        return self._run_role_phase("agent2", state, config)

    def agent3_competitor_perspective(
        self,
        state: DeepReviewState,
        runtime: Runtime[AgentContext] | None = None,
        config: RunnableConfig | None = None,
    ) -> DeepReviewState:
        """Agent 3: autonomous competitor/stakeholder simulator."""
        return self._run_role_phase("agent3", state, config)

    def increment_iteration(self, state: DeepReviewState) -> DeepReviewState:
        current_iteration = state["current_iteration"] + 1
        return {"current_iteration": current_iteration}

    def generate_solution(
        self,
        state: DeepReviewState,
        config: RunnableConfig | None = None,
    ) -> DeepReviewState:
        """Generate the overall, refined solution based on all iterations."""
        events = self.events(config)
        events.emit(
            "Synthesizing final hypothesis",
            stage="finalize",
            iteration=state["current_iteration"],
        )
        prompt = f"Original question: {state['question']}\n\n"
        prompt += "Evolution of solutions:\n"

        for i, (solution_text, critique_text, perspective_text) in enumerate(
            zip(
                state["agent1_solution"],
                state["agent2_critiques"],
                state["agent3_perspectives"],
            ),
            start=1,
        ):
            prompt += f"\nIteration {i}:\n"
            prompt += f"Solution: {solution_text}\n"
            prompt += f"Critique: {critique_text}\n"
            prompt += f"Competitor perspective: {perspective_text}\n"

        prompt += "\nBased on this iterative process, provide the overall, refined solution."
        solution = self.strllm.invoke(prompt)
        events.emit(
            "Final hypothesis ready",
            stage="finalize_result",
            preview=solution,
        )
        return {"solution": solution}

    def print_visited_sites(
        self,
        state: DeepReviewState,
        config: RunnableConfig | None = None,
    ) -> DeepReviewState:
        return state.copy()

    @staticmethod
    def _escape_latex(text: Any) -> str:
        replacements = {
            "\\": r"\textbackslash{}",
            "&": r"\&",
            "%": r"\%",
            "$": r"\$",
            "#": r"\#",
            "_": r"\_",
            "{": r"\{",
            "}": r"\}",
            "~": r"\textasciitilde{}",
            "^": r"\textasciicircum{}",
        }
        escaped = str(text or "")
        for char, replacement in replacements.items():
            escaped = escaped.replace(char, replacement)
        return escaped

    @staticmethod
    def _strip_latex_fences(text: str) -> str:
        stripped = text.strip()
        if stripped.startswith("```latex"):
            stripped = stripped.removeprefix("```latex").strip()
        elif stripped.startswith("```tex"):
            stripped = stripped.removeprefix("```tex").strip()
        elif stripped.startswith("```"):
            stripped = stripped.removeprefix("```").strip()
        if stripped.endswith("```"):
            stripped = stripped.removesuffix("```").strip()
        return stripped

    def _format_websites_latex(self, visited_sites: set[str] | None) -> str:
        if not visited_sites:
            return ""
        items = "\n".join(
            f"  \\item \\url{{{self._escape_latex(site)}}}"
            for site in sorted(visited_sites)
        )
        return (
            "\\section*{Websites Used in Research}\n"
            "\\begin{itemize}\n"
            f"{items}\n"
            "\\end{itemize}\n"
        )

    def _ensure_complete_latex_document(
        self,
        latex_response: str,
        state: DeepReviewState,
    ) -> str:
        latex_doc = self._strip_latex_fences(latex_response)
        if r"\documentclass" in latex_doc:
            if r"\end{document}" not in latex_doc:
                latex_doc = latex_doc.rstrip() + "\n\\end{document}"
            return latex_doc

        solution = self._escape_latex(state.get("solution", ""))
        question = self._escape_latex(state.get("question", ""))
        model_response = self._escape_latex(latex_response)
        return f"""\\documentclass{{article}}
\\usepackage[margin=1in]{{geometry}}
\\usepackage{{hyperref}}
\\begin{{document}}
\\section*{{Executive Summary}}
This report summarizes the deep-review process for the following question:

{question}

\\section*{{Final Solution}}
{solution}

\\section*{{Raw Model Report}}
{model_response}

\\end{{document}}"""

    @staticmethod
    def _inject_before_end_document(original_tex: str, injection: str) -> str:
        if not injection.strip():
            return original_tex
        injection_index = original_tex.rfind(r"\end{document}")
        if injection_index == -1:
            return original_tex.rstrip() + "\n" + injection
        return (
            original_tex[:injection_index]
            + "\n"
            + injection
            + "\n"
            + original_tex[injection_index:]
        )

    def summarize_process_as_latex(
        self,
        state: DeepReviewState,
        config: RunnableConfig | None = None,
    ) -> DeepReviewState:
        """Summarize the iterative review process as a valid LaTeX document."""
        events = self.events(config)
        events.emit(
            "Writing LaTeX report",
            stage="summarize",
        )
        iteration_details = ""
        for i, (sol, crit, comp) in enumerate(
            zip(
                state["agent1_solution"],
                state["agent2_critiques"],
                state["agent3_perspectives"],
            ),
            start=1,
        ):
            iteration_details += (
                f"\\subsection*{{Iteration {i}}}\n\n"
                f"\\textbf{{Solution:}}\\\\\n{sol}\n\n"
                f"\\textbf{{Critique:}}\\\\\n{crit}\n\n"
                f"\\textbf{{Competitor Perspective:}}\\\\\n{comp}\n\n"
            )

        timestamp_str = datetime.now(UTC).strftime("%Y-%m-%d_%H-%M-%S")
        txt_filename = Path(
            self.den,
            f"iteration_details_{timestamp_str}_chat_history.txt",
        )
        with open(txt_filename, "w", encoding="utf-8") as f:
            f.write(iteration_details)

        prompt = f"""\
            You are a system that produces a FULL LaTeX document.
            Here is information about a multi-iteration process:

            Original question: {state["question"]}

            Below are the solutions, critiques, and competitor perspectives from each iteration:

            {iteration_details}

            The solution we arrived at was:

            {state["solution"]}

            Now produce a valid LaTeX document. Be sure to use a table of contents.
            It must start with an Executive Summary (that may be multiple pages) which summarizes
            the entire iterative process. Following that, include the solution in full,
            not summarized, but reformatted for appropriate LaTeX. Finally, include all
            steps - solutions, critiques, and competitor perspectives - in an Appendix,
            and include a listing of all websites used in research if any were used.

            You must ONLY RETURN LaTeX, nothing else. It must be valid LaTeX syntax!

            Your output should start with:
            \\documentclass{{article}}
            \\usepackage[margin=1in]{{geometry}}
            etc.

            It must compile without errors under pdflatex.
        """

        websites_latex = self._format_websites_latex(state.get("visited_sites"))
        latex_response = self.strllm.invoke(prompt)
        latex_doc = self._ensure_complete_latex_document(latex_response, state)
        final_latex = self._inject_before_end_document(
            latex_doc, websites_latex
        )
        events.emit(
            "Report ready",
            stage="summarize_result",
            preview=latex_response,
            output_path=str(txt_filename),
        )
        return {"summary_report": final_latex}

    def _build_graph(self):
        self.tool_llm = self.tool_llm.bind_tools(self.tools.values())

        self.add_node(self.agent1_generate_solution, "agent1")
        self.add_node(self.agent2_critique, "agent2")
        self.add_node(self.agent3_competitor_perspective, "agent3")
        self.add_node(self.tool_node, "tool_node")
        self.add_node(self.increment_iteration, "increment_iteration")
        self.add_node(self.generate_solution, "finalize")
        self.add_node(self.print_visited_sites, "print_sites")
        self.add_node(self.summarize_process_as_latex, "summarize_as_latex")

        self.graph.set_entry_point("agent1")

        self.graph.add_conditional_edges(
            "agent1",
            self._wrap_cond(
                phase_should_continue, "agent1_continue", "deep_review"
            ),
            {"tools": "tool_node", "done": "agent2"},
        )
        self.graph.add_conditional_edges(
            "agent2",
            self._wrap_cond(
                phase_should_continue, "agent2_continue", "deep_review"
            ),
            {"tools": "tool_node", "done": "agent3"},
        )
        self.graph.add_conditional_edges(
            "agent3",
            self._wrap_cond(
                phase_should_continue, "agent3_continue", "deep_review"
            ),
            {"tools": "tool_node", "done": "increment_iteration"},
        )
        self.graph.add_conditional_edges(
            "tool_node",
            self._wrap_cond(
                route_after_tools, "route_after_tools", "deep_review"
            ),
            {"agent1": "agent1", "agent2": "agent2", "agent3": "agent3"},
        )

        self.graph.add_conditional_edges(
            "increment_iteration",
            self._wrap_cond(should_continue, "should_continue", "deep_review"),
            {"continue": "agent1", "finish": "finalize"},
        )

        self.graph.add_edge("finalize", "summarize_as_latex")
        self.graph.add_edge("summarize_as_latex", "print_sites")
        self.graph.set_finish_point("print_sites")

agent1_generate_solution(state, runtime=None, config=None)

Agent 1: autonomous solution generator.

Source code in src/ursa/agents/deep_review_agent.py

def agent1_generate_solution(
    self,
    state: DeepReviewState,
    runtime: Runtime[AgentContext] | None = None,
    config: RunnableConfig | None = None,
) -> DeepReviewState:
    """Agent 1: autonomous solution generator."""
    return self._run_role_phase("agent1", state, config)

agent2_critique(state, runtime=None, config=None)

Agent 2: autonomous critic.

Source code in src/ursa/agents/deep_review_agent.py

def agent2_critique(
    self,
    state: DeepReviewState,
    runtime: Runtime[AgentContext] | None = None,
    config: RunnableConfig | None = None,
) -> DeepReviewState:
    """Agent 2: autonomous critic."""
    return self._run_role_phase("agent2", state, config)

agent3_competitor_perspective(state, runtime=None, config=None)

Agent 3: autonomous competitor/stakeholder simulator.

Source code in src/ursa/agents/deep_review_agent.py

def agent3_competitor_perspective(
    self,
    state: DeepReviewState,
    runtime: Runtime[AgentContext] | None = None,
    config: RunnableConfig | None = None,
) -> DeepReviewState:
    """Agent 3: autonomous competitor/stakeholder simulator."""
    return self._run_role_phase("agent3", state, config)

generate_solution(state, config=None)

Generate the overall, refined solution based on all iterations.

Source code in src/ursa/agents/deep_review_agent.py

def generate_solution(
    self,
    state: DeepReviewState,
    config: RunnableConfig | None = None,
) -> DeepReviewState:
    """Generate the overall, refined solution based on all iterations."""
    events = self.events(config)
    events.emit(
        "Synthesizing final hypothesis",
        stage="finalize",
        iteration=state["current_iteration"],
    )
    prompt = f"Original question: {state['question']}\n\n"
    prompt += "Evolution of solutions:\n"

    for i, (solution_text, critique_text, perspective_text) in enumerate(
        zip(
            state["agent1_solution"],
            state["agent2_critiques"],
            state["agent3_perspectives"],
        ),
        start=1,
    ):
        prompt += f"\nIteration {i}:\n"
        prompt += f"Solution: {solution_text}\n"
        prompt += f"Critique: {critique_text}\n"
        prompt += f"Competitor perspective: {perspective_text}\n"

    prompt += "\nBased on this iterative process, provide the overall, refined solution."
    solution = self.strllm.invoke(prompt)
    events.emit(
        "Final hypothesis ready",
        stage="finalize_result",
        preview=solution,
    )
    return {"solution": solution}

summarize_process_as_latex(state, config=None)

Summarize the iterative review process as a valid LaTeX document.

Source code in src/ursa/agents/deep_review_agent.py

def summarize_process_as_latex(
    self,
    state: DeepReviewState,
    config: RunnableConfig | None = None,
) -> DeepReviewState:
    """Summarize the iterative review process as a valid LaTeX document."""
    events = self.events(config)
    events.emit(
        "Writing LaTeX report",
        stage="summarize",
    )
    iteration_details = ""
    for i, (sol, crit, comp) in enumerate(
        zip(
            state["agent1_solution"],
            state["agent2_critiques"],
            state["agent3_perspectives"],
        ),
        start=1,
    ):
        iteration_details += (
            f"\\subsection*{{Iteration {i}}}\n\n"
            f"\\textbf{{Solution:}}\\\\\n{sol}\n\n"
            f"\\textbf{{Critique:}}\\\\\n{crit}\n\n"
            f"\\textbf{{Competitor Perspective:}}\\\\\n{comp}\n\n"
        )

    timestamp_str = datetime.now(UTC).strftime("%Y-%m-%d_%H-%M-%S")
    txt_filename = Path(
        self.den,
        f"iteration_details_{timestamp_str}_chat_history.txt",
    )
    with open(txt_filename, "w", encoding="utf-8") as f:
        f.write(iteration_details)

    prompt = f"""\
        You are a system that produces a FULL LaTeX document.
        Here is information about a multi-iteration process:

        Original question: {state["question"]}

        Below are the solutions, critiques, and competitor perspectives from each iteration:

        {iteration_details}

        The solution we arrived at was:

        {state["solution"]}

        Now produce a valid LaTeX document. Be sure to use a table of contents.
        It must start with an Executive Summary (that may be multiple pages) which summarizes
        the entire iterative process. Following that, include the solution in full,
        not summarized, but reformatted for appropriate LaTeX. Finally, include all
        steps - solutions, critiques, and competitor perspectives - in an Appendix,
        and include a listing of all websites used in research if any were used.

        You must ONLY RETURN LaTeX, nothing else. It must be valid LaTeX syntax!

        Your output should start with:
        \\documentclass{{article}}
        \\usepackage[margin=1in]{{geometry}}
        etc.

        It must compile without errors under pdflatex.
    """

    websites_latex = self._format_websites_latex(state.get("visited_sites"))
    latex_response = self.strllm.invoke(prompt)
    latex_doc = self._ensure_complete_latex_document(latex_response, state)
    final_latex = self._inject_before_end_document(
        latex_doc, websites_latex
    )
    events.emit(
        "Report ready",
        stage="summarize_result",
        preview=latex_response,
        output_path=str(txt_filename),
    )
    return {"summary_report": final_latex}

dsi_agent

DSIAgent

Bases: AgentWithTools, BaseAgent[DSIState]

Source code in src/ursa/agents/dsi_agent.py

class DSIAgent(AgentWithTools, BaseAgent[DSIState]):
    state_type = DSIState

    def __init__(
        self,
        llm: BaseChatModel,
        database_path: str = "",
        output_mode: str = "agent",
        run_path: str = "",
        thread_id: str | None = None,
        **kwargs,
    ):
        default_tools = [
            run_web_search,
            run_osti_search,
            run_arxiv_search,
            load_dsi_tool,
            query_dsi_tool,
            download_file_tool,
            read_file,
            run_command,
            write_code,
            edit_code,
        ]

        super().__init__(llm=llm, tools=default_tools, **kwargs)
        self.db_schema = ""
        self.db_description = ""

        self.master_db_folder = ""
        self.master_database_path = ""
        self.current_db_abs_path = ""

        self.output_mode = output_mode

        # Try to load the master database if a path is provided, otherwise wait for the user to load one
        if run_path == "":
            self.run_path = os.getcwd()
        else:
            self.run_path = run_path
        self.load_master_db(database_path)

        self.prompt = """
        You are a data-analysis agent who can write python code, SQL queries, and generate plots to answer user questions based on the data available in a DSI object.
        Use the load_dsi_tool tool to load DSI files that have a .db extension
        Use query_dsi_tool to run SQL queries on it
        When a user asks for data or dataset or ... you have, do NOT list the schema or metadata information you have about tables. Query the DSI objects for data and list the data in the tables.

        You can:
        - write and execute Python code,
        - compose SQL statements but ONLY select; no update or delete
        - generate plots and diagrams,
        - analyze and summarize data.

        Requirements:
        - Planning: Think carefully about the problem, but **do not show your reasoning**.
        - Data:
            - Always use the provided tools when available — never simulate results.
            - Never fabricate or assume sample data. Query or compute real values only.
            - When creating plots or files, always save them directly to disk (do not embed inline output).
            - Do not infer or assume any data beyond what is provided by the tools.
        - Keep all responses concise and focused on the requested task.
        - Only load a dataset when explicitly asked by a user
        - Do not restate the prompt or reasoning; just act and report the outcome briefly.
        """

        if thread_id:
            self.thread_id = thread_id
        else:
            self.thread_id = str(random.randint(1, 20000))

    def load_master_db(self, master_database: str) -> None:
        """Load the  master dataset from the given path.

        Arg:
            master_database (str): the path to the DSI object
        """

        if master_database == "":
            LOGGER.error("No DSI database provided. Please load one")
            return

        _master_database_path, _master_db_folder = _get_db_abs_path(
            master_database, self.run_path
        )
        absolute_db_path = _master_database_path

        if check_db_valid(absolute_db_path):
            self.db_tables, self.db_schema, self.db_description = get_db_info(
                absolute_db_path
            )

            # set the values now that we know things are correct
            self.current_db_abs_path = absolute_db_path
            self.master_database_path = absolute_db_path
            self.master_db_folder = _master_db_folder

        else:
            LOGGER.error("No valid DSI database provided. Please load one")
            # sys.exit(1)

    # __call__ from my agent
    def _response_node(self, state):
        new_state, full_overwrite = self.prepare_messages_context(state)
        messages = new_state["messages"]

        conversation = [SystemMessage(content=self.prompt)] + messages
        response = self.llm.invoke(conversation)

        return self.messages_update(
            new_state,
            [response],
            full_overwrite=full_overwrite,
            extra={
                "response": response.content,
                "metadata": response.response_metadata,
            },
        )

    def _build_graph(self):
        self.llm = self.llm.bind_tools(self.tools.values())

        self.add_node(self._response_node, "response")
        self.add_node(self.tool_node, "tools")

        self.graph.add_edge(START, "response")
        self.graph.add_conditional_edges(
            "response",
            self._wrap_cond(should_call_tools, "should_call_tools", "dsi"),
            {
                "call_tools": "tools",
                "continue": END,
            },
        )
        self.graph.add_edge("tools", "response")

    def craft_message(self, human_msg):
        """Craft the message with context if available."""

        base_system_context = f"""
            The following phrases all refer to this same database:
            - "master database"
            - "master dataset"
            - "DSIExplorer master database"
            - "Diana dataset

            When the user asks to reload, refresh, reset, reinitialize, restart, or 
            the **master database**, interpret that as a request to reload the 
            DSIExplorer master database using the tool load_dsi_tool("{self.master_database_path}"),
            load the last dataset in the context.

            Do no reload or load the master dataset unless explicitly asked by the user.
        """

        # Build remaining dynamic system context parts
        system_parts = [base_system_context]

        if self.current_db_abs_path != "":
            system_parts.append(
                "The current working database path (current_db_abs_path) is: "
                + self.current_db_abs_path
            )

        if self.master_database_path != "":
            system_parts.append(
                "The master database path (master_database_path) is: "
                + self.master_database_path
            )

        if self.db_schema != "":
            system_parts.append(
                "The schema of the dataset loaded: " + self.db_schema
            )

        if self.db_description != "":
            system_parts.append("Dataset description: " + self.db_description)

        # Combine
        if system_parts:
            system_message = SystemMessage(content="\n\n".join(system_parts))
            messages = [system_message, HumanMessage(content=human_msg)]
        else:
            messages = [HumanMessage(content=human_msg)]

        # clear
        self.db_schema = ""
        self.db_description = ""
        self.master_database_path = ""
        self.current_db_abs_path = ""

        return {"messages": messages}

    def format_query(
        self, user_query, state: DSIState | None = None
    ) -> DSIState:
        """
           Injest string query into the agent state

        Arg:
           user_query (str): the user question
        """
        if state is not None and "messages" in state:
            pass  # This is where we should process it if we want the state
            # passed in to be appended to. Not sure if we do so reverting
            # to current behavior for now
            # state["messages"].append(HumanMessage(content=str(user_query)))
            # return state
        return self.craft_message(user_query)

    def format_result(self, result: DSIState, start_time=None) -> str:
        """
        Parse result state into the desired output string

        Arg:
           result (DSIState): The state output from the DSI agent
        """
        response_text = result["response"]
        cleaned_output = response_text.strip()

        return cleaned_output

    def ask(self, user_query) -> None:
        """Ask a question to the DSI Explorer agent.

        Arg:
            user_query (str): the user question
        """

        start = now()

        msg = self.format_query(user_query)

        result = self.invoke(
            msg, config={"configurable": {"thread_id": self.thread_id}}
        )

        formatted_result = self.format_result(result, start)

        # I would like to add this
        # if self.output_mode == "jupyter":
        #     from IPython.display import Markdown, display
        #     display(Markdown(formatted_result))

        return formatted_result

ask(user_query)

Ask a question to the DSI Explorer agent.

Arg

user_query (str): the user question

Source code in src/ursa/agents/dsi_agent.py

def ask(self, user_query) -> None:
    """Ask a question to the DSI Explorer agent.

    Arg:
        user_query (str): the user question
    """

    start = now()

    msg = self.format_query(user_query)

    result = self.invoke(
        msg, config={"configurable": {"thread_id": self.thread_id}}
    )

    formatted_result = self.format_result(result, start)

    # I would like to add this
    # if self.output_mode == "jupyter":
    #     from IPython.display import Markdown, display
    #     display(Markdown(formatted_result))

    return formatted_result

craft_message(human_msg)

Craft the message with context if available.

Source code in src/ursa/agents/dsi_agent.py

def craft_message(self, human_msg):
    """Craft the message with context if available."""

    base_system_context = f"""
        The following phrases all refer to this same database:
        - "master database"
        - "master dataset"
        - "DSIExplorer master database"
        - "Diana dataset

        When the user asks to reload, refresh, reset, reinitialize, restart, or 
        the **master database**, interpret that as a request to reload the 
        DSIExplorer master database using the tool load_dsi_tool("{self.master_database_path}"),
        load the last dataset in the context.

        Do no reload or load the master dataset unless explicitly asked by the user.
    """

    # Build remaining dynamic system context parts
    system_parts = [base_system_context]

    if self.current_db_abs_path != "":
        system_parts.append(
            "The current working database path (current_db_abs_path) is: "
            + self.current_db_abs_path
        )

    if self.master_database_path != "":
        system_parts.append(
            "The master database path (master_database_path) is: "
            + self.master_database_path
        )

    if self.db_schema != "":
        system_parts.append(
            "The schema of the dataset loaded: " + self.db_schema
        )

    if self.db_description != "":
        system_parts.append("Dataset description: " + self.db_description)

    # Combine
    if system_parts:
        system_message = SystemMessage(content="\n\n".join(system_parts))
        messages = [system_message, HumanMessage(content=human_msg)]
    else:
        messages = [HumanMessage(content=human_msg)]

    # clear
    self.db_schema = ""
    self.db_description = ""
    self.master_database_path = ""
    self.current_db_abs_path = ""

    return {"messages": messages}

format_query(user_query, state=None)

Injest string query into the agent state

Arg

user_query (str): the user question

Source code in src/ursa/agents/dsi_agent.py

def format_query(
    self, user_query, state: DSIState | None = None
) -> DSIState:
    """
       Injest string query into the agent state

    Arg:
       user_query (str): the user question
    """
    if state is not None and "messages" in state:
        pass  # This is where we should process it if we want the state
        # passed in to be appended to. Not sure if we do so reverting
        # to current behavior for now
        # state["messages"].append(HumanMessage(content=str(user_query)))
        # return state
    return self.craft_message(user_query)

format_result(result, start_time=None)

Parse result state into the desired output string

Arg

result (DSIState): The state output from the DSI agent

Source code in src/ursa/agents/dsi_agent.py

def format_result(self, result: DSIState, start_time=None) -> str:
    """
    Parse result state into the desired output string

    Arg:
       result (DSIState): The state output from the DSI agent
    """
    response_text = result["response"]
    cleaned_output = response_text.strip()

    return cleaned_output

load_master_db(master_database)

Load the master dataset from the given path.

Arg

master_database (str): the path to the DSI object

Source code in src/ursa/agents/dsi_agent.py

def load_master_db(self, master_database: str) -> None:
    """Load the  master dataset from the given path.

    Arg:
        master_database (str): the path to the DSI object
    """

    if master_database == "":
        LOGGER.error("No DSI database provided. Please load one")
        return

    _master_database_path, _master_db_folder = _get_db_abs_path(
        master_database, self.run_path
    )
    absolute_db_path = _master_database_path

    if check_db_valid(absolute_db_path):
        self.db_tables, self.db_schema, self.db_description = get_db_info(
            absolute_db_path
        )

        # set the values now that we know things are correct
        self.current_db_abs_path = absolute_db_path
        self.master_database_path = absolute_db_path
        self.master_db_folder = _master_db_folder

    else:
        LOGGER.error("No valid DSI database provided. Please load one")

check_db_valid(db_path)

Check if the provided path points to a valid DSI database.

Arg

db_path (str): the absolute path of the DSI database

Returns:

Name	Type	Description
bool	bool	True if the database is valid, False otherwise

Source code in src/ursa/agents/dsi_agent.py

def check_db_valid(db_path: str) -> bool:
    """Check if the provided path points to a valid DSI database.

    Arg:
        db_path (str): the absolute path of the DSI database

    Returns:
        bool: True if the database is valid, False otherwise
    """
    if not os.path.exists(db_path):
        return False
    else:
        try:
            with redirect_stdout(_NULL), redirect_stderr(_NULL):
                temp_store = DSI(db_path, check_same_thread=False)
                temp_store.list(
                    True
                )  # force things to fail if the table is empty
                temp_store.close()

        except Exception:  # noqa: BLE001
            return False

    return True

get_db_info(db_path)

Load the database information (tables and schema) from a DSI database.

Arg

db_path (str): the absolute path of the DSI database

Returns:

Name	Type	Description
list	list	the list of tables in the database
dict	dict	the schema of the database
str	str	the description of the database (if available, otherwise empty string)

Source code in src/ursa/agents/dsi_agent.py

def get_db_info(db_path: str) -> tuple[list, dict, str]:
    """Load the database information (tables and schema) from a DSI database.

    Arg:
        db_path (str): the absolute path of the DSI database

    Returns:
        list: the list of tables in the database
        dict: the schema of the database
        str: the description of the database (if available, otherwise empty string)
    """

    tables = []
    schema = {}
    desc = ""

    if check_db_valid(db_path) is False:
        return tables, schema, desc

    try:
        # with open(os.devnull, "w") as fnull:
        #     with redirect_stdout(fnull), redirect_stderr(fnull):
        with redirect_stdout(_NULL), redirect_stderr(_NULL):
            _dsi_store = DSI(db_path, check_same_thread=False)
            tables = _dsi_store.list(True)
            schema = _dsi_store.schema()
            desc = load_db_description(db_path)
            _dsi_store.close()

        return tables, schema, desc

    except Exception:  # noqa: BLE001
        return tables, schema, desc

load_db_description(db_path)

Load the database description from a YAML file when provided with the path to a DSI database.

Arg

db_path (str): the absolute path of the DSI database

Returns:

Name	Type	Description
str	str	message indicating success or failure

Source code in src/ursa/agents/dsi_agent.py

def load_db_description(db_path: str) -> str:
    """Load the database description from a YAML file when provided with the path to a DSI database.

    Arg:
        db_path (str): the absolute path of the DSI database

    Returns:
        str: message indicating success or failure
    """

    try:
        # The description file is expected to be in the same directory as the database, with the same name but ending in '_description.yaml'
        db_description_path = db_path.rsplit(".", 1)[0] + "_description.yaml"

        with open(db_description_path, "r") as f:
            db_desc = f.read()

        return str(db_desc)
    except Exception:  # noqa: BLE001
        return ""

should_call_tools(state)

Decide whether to call tools or continue.

Arg

state (State): the current state of the graph

Returns:

Name	Type	Description
str	str	"call_tools" or "continue"

Source code in src/ursa/agents/dsi_agent.py

def should_call_tools(state: DSIState) -> str:
    """Decide whether to call tools or continue.

    Arg:
        state (State): the current state of the graph

    Returns:
        str: "call_tools" or "continue"
    """

    last = state["messages"][-1]
    if isinstance(last, AIMessage) and last.tool_calls:
        return "call_tools"

    return "continue"

execution_agent

Execution agent that builds a tool-enabled state graph to autonomously run tasks.

This module implements ExecutionAgent, a LangGraph-based agent that executes user instructions by invoking LLM tool calls and coordinating a controlled workflow.

Key features: - Workspace management with optional symlinking for external sources. - Safety-checked shell execution via run_command with output size budgeting. - Code authoring and edits through write_code and edit_code. - Web search capability through DuckDuckGoSearchResults. - Summarization of the session - Configurable graph with nodes for agent, action, and summarize.

Implementation notes: - LLM prompts are sourced from prompt_library.execution_prompts. - Outputs from subprocess are trimmed under MAX_TOOL_MSG_CHARS to fit tool messages. - The agent uses ToolNode and LangGraph StateGraph to loop until no tool calls remain. - Safety gates block unsafe shell commands and surface the rationale to the user.

Environment: - MAX_TOOL_MSG_CHARS caps combined stdout/stderr in tool responses.

Entry points: - ExecutionAgent._invoke(...) runs the compiled graph. - main() shows a minimal demo that writes and runs a script.

ExecutionAgent

Bases: AgentWithTools, BaseAgent[ExecutionState]

Orchestrates model-driven code execution, tool calls, and state management.

Orchestrates model-driven code execution, tool calls, and state management for iterative program synthesis and shell interaction.

This agent wraps an LLM with a small execution graph that alternates between issuing model queries, invoking tools (read, run, write, edit, search), performing safety checks, and summarizing progress. It manages a workspace on disk, optional symlinks.

Parameters:

Name	Type	Description	Default
llm	BaseChatModel	Model identifier or bound chat model instance. If a string is provided, the BaseAgent initializer will resolve it.	required
log_state	bool	When True, the agent writes intermediate json state to disk for debugging and auditability.	False
**kwargs		Passed through to the BaseAgent constructor (e.g., model configuration, checkpointer).	{}

Attributes:

Name	Type	Description
safe_codes	list[str]	List of trusted programming languages for the agent. Defaults to python and julia
executor_prompt	str	Prompt used when invoking the executor LLM loop.
recap_prompt	str	Prompt used to request concise summaries for final output.
tools	dict[str, Tool]	Tools available to the agent (run_command, write_code, edit_code, read_file, run_web_search, run_osti_search, run_arxiv_search), keyed by tool name for quick lookups.
tool_node	ToolNode	Graph node that dispatches tool calls.
llm	BaseChatModel	Base LLM instance used for summaries and recap.
tool_llm	BaseChatModel	Tool-bound LLM used for the execution loop.

Methods:

Name	Description
query_executor	Send messages to the executor LLM, ensure workspace exists, and handle symlink setup before returning the model response.
recap	Produce a summary of recent interactions.
_build_graph	Construct and compile the StateGraph for the agent loop.

Raises:

Type	Description
AttributeError	Accessing the .action attribute raises to encourage using .stream(...) or .invoke(...).

Source code in src/ursa/agents/execution_agent.py

class ExecutionAgent(AgentWithTools, BaseAgent[ExecutionState]):
    """Orchestrates model-driven code execution, tool calls, and state management.

    Orchestrates model-driven code execution, tool calls, and state management for
    iterative program synthesis and shell interaction.

    This agent wraps an LLM with a small execution graph that alternates
    between issuing model queries, invoking tools (read, run, write, edit, search),
    performing safety checks, and summarizing progress. It manages a
    workspace on disk, optional symlinks.

    Args:
        llm (BaseChatModel): Model identifier or bound chat model
            instance. If a string is provided, the BaseAgent initializer will
            resolve it.
        log_state (bool): When True, the agent writes intermediate json state
            to disk for debugging and auditability.
        **kwargs: Passed through to the BaseAgent constructor (e.g., model
            configuration, checkpointer).

    Attributes:
        safe_codes (list[str]): List of trusted programming languages for the
            agent. Defaults to python and julia
        executor_prompt (str): Prompt used when invoking the executor LLM
            loop.
        recap_prompt (str): Prompt used to request concise summaries for
            final output.
        tools (dict[str, Tool]): Tools available to the agent (run_command, write_code,
            edit_code, read_file, run_web_search, run_osti_search, run_arxiv_search),
            keyed by tool name for quick lookups.
        tool_node (ToolNode): Graph node that dispatches tool calls.
        llm (BaseChatModel): Base LLM instance used for summaries and recap.
        tool_llm (BaseChatModel): Tool-bound LLM used for the execution loop.

    Methods:
        query_executor(state): Send messages to the executor LLM, ensure
            workspace exists, and handle symlink setup before returning the
            model response.
        recap(state): Produce a summary of recent interactions.
        _build_graph(): Construct and compile the StateGraph for the agent
            loop.

    Raises:
        AttributeError: Accessing the .action attribute raises to encourage
            using .stream(...) or .invoke(...).
    """

    state_type = ExecutionState

    def __init__(
        self,
        llm: BaseChatModel,
        log_state: bool = False,
        extra_tools: list[BaseTool] | None = None,
        tokens_before_summarize: int = 50000,
        messages_to_keep: int = 20,
        use_web: bool = False,
        safe_codes: list[str] | None = None,
        **kwargs,
    ):
        default_tools = [
            run_command,
            write_code,
            edit_code,
            read_file,
            download_file_tool,
            read_image_tool,
            read_experience,
            write_experience,
            edit_experience,
            list_experiences,
        ]
        if use_web:
            default_tools.extend([
                run_web_search,
                run_osti_search,
                run_arxiv_search,
            ])
        if extra_tools:
            default_tools.extend(extra_tools)

        super().__init__(
            llm=llm,
            tools=default_tools,
            safe_codes=safe_codes or ["python", "julia"],
            **kwargs,
        )
        self.executor_prompt = executor_prompt
        self.recap_prompt = recap_prompt
        self.extra_tools = extra_tools
        self.log_state = log_state
        self.tokens_before_summarize = tokens_before_summarize
        self.messages_to_keep = messages_to_keep

    @staticmethod
    def _message_text(msg: Any) -> str:
        """Extract readable text from LangChain-style message inputs."""
        text = getattr(msg, "text", None)
        if text:
            return str(text)
        content = getattr(msg, "content", None)
        if content is not None:
            return str(content)
        if isinstance(msg, Mapping):
            content = msg.get("content")
            if content is not None:
                return str(content)
        if isinstance(msg, tuple) and len(msg) >= 2:
            return str(msg[1])
        return ""

    @classmethod
    def _first_human_message_text(cls, messages: Any) -> str:
        """Return text for the first human/user message in a message sequence."""
        if not isinstance(messages, (list, tuple)):
            return ""
        for msg in messages:
            if isinstance(msg, HumanMessage):
                return cls._message_text(msg)
            if isinstance(msg, Mapping):
                role = str(msg.get("role") or msg.get("type") or "").lower()
                if role in {"human", "user"}:
                    return cls._message_text(msg)
            if isinstance(msg, tuple) and msg:
                role = str(msg[0]).lower()
                if role in {"human", "user"}:
                    return cls._message_text(msg)
        return ""

    def _normalize_inputs(self, inputs):
        """Normalize inputs and remember this invocation's user request.

        Persistent LangGraph threads merge a new ``messages`` input with prior
        checkpointed history. Capturing the first human/user message before that
        merge gives review a stable request for the current invoke instead of the
        first human message in the entire persisted conversation.
        """
        normalized = dict(super()._normalize_inputs(inputs))
        user_request = self._first_human_message_text(
            normalized.get("messages", [])
        )
        normalized["current_user_request"] = user_request
        return normalized

    # Define the function that calls the model
    def query_executor(
        self,
        state: ExecutionState,
        runtime: Runtime[AgentContext],
        config: RunnableConfig | None = None,
    ) -> ExecutionState:
        """Prepare workspace, handle optional symlinks, and invoke the executor LLM.

        This method copies the incoming state, ensures a workspace directory exists
        (creating one with a default name when absent), optionally creates a symlink
        described by state["symlinkdir"], sets or injects the executor system prompt
        as the first message, and invokes the bound LLM. When logging is enabled,
        it persists the pre-invocation state to disk.

        Args:
            state: The current execution state. Expected keys include:
                - "messages": Ordered list of System/Human/AI/Tool messages.
                - "workspace": Optional path to the working directory.
                - "symlinkdir": Optional dict with "source" and "dest" keys.

        Returns:
            ExecutionState: Partial state update containing:
                - "messages": A list with the model's response as the latest entry.
                - "workspace": The resolved workspace path.
        """
        # Add model to the state so it can be passed to tools like the URSA Arxiv or OSTI tools
        new_state = deepcopy(state)
        events = self.events(config)
        new_state.setdefault("symlinkdir", {})

        full_overwrite = False

        # 1.5) Check message history length, summarize, and patch dangling tool calls.
        new_state, full_overwrite = self.prepare_messages_context(
            new_state, system_prompt=self.executor_prompt
        )

        # 2) Optionally create a symlink if symlinkdir is provided and not yet linked.
        sd = new_state.get("symlinkdir")
        if sd and "is_linked" not in sd:
            # symlinkdir structure: {"source": "/path/to/src", "dest": "link/name"}
            symlinkdir = sd

            src = Path(symlinkdir["source"]).expanduser().resolve()
            dst = runtime.context.workspace.joinpath(symlinkdir["dest"])

            # If a file/link already exists at the destination, replace it.
            if dst.exists() or dst.is_symlink():
                dst.unlink()

            # Ensure parent directories for the link exist.
            dst.parent.mkdir(parents=True, exist_ok=True)

            # Create the symlink (tell pathlib if the target is a directory).
            dst.symlink_to(src, target_is_directory=src.is_dir())
            events.emit(
                "Workspace symlink created",
                stage="workspace_symlink",
                source=str(src),
                destination=str(dst),
            )
            new_state["symlinkdir"]["is_linked"] = True
            full_overwrite = True

        # 3) Ensure the executor prompt is the first SystemMessage.
        messages = deepcopy(new_state["messages"])
        if isinstance(messages[0], SystemMessage):
            messages[0] = SystemMessage(content=self.executor_prompt)
        else:
            messages = [SystemMessage(content=self.executor_prompt)] + messages

        # 4) Invoke the LLM with the prepared message sequence.
        try:
            response = self.tool_llm.invoke(
                messages, self.build_config(tags=["agent"])
            )
            new_state["messages"].append(response)
        except Exception as e:  # noqa: BLE001
            response = AIMessage(content=f"Response error {e}")
            msg = new_state["messages"][-1].text
            events.emit(
                "Executor response failed",
                stage="executor_error",
                error_type=type(e).__name__,
                error=str(e),
                latest_message=msg[:2000],
            )
            new_state["messages"].append(response)

        # 5) Optionally persist the pre-invocation state for audit/debugging.
        if self.log_state:
            self.write_state("execution_agent.json", new_state)
        if full_overwrite:
            return {
                "messages": Overwrite(new_state["messages"]),
                "symlinkdir": new_state["symlinkdir"],
            }
        else:
            return {"messages": response, "symlinkdir": new_state["symlinkdir"]}

    def _get_current_user_request(self, state: ExecutionState) -> str:
        """Return the user request for the current invoke.

        Prefer ``current_user_request``, which is captured from the entry input
        before LangGraph merges new messages with any checkpointed history. If a
        legacy/in-memory state does not have that field, fall back to the latest
        human message rather than the first one in the full history.
        """
        current_request = state.get("current_user_request")
        if current_request:
            return str(current_request)
        for msg in reversed(state.get("messages", [])):
            if isinstance(msg, HumanMessage):
                return self._message_text(msg)
        messages = state.get("messages", [])
        if messages:
            return self._message_text(messages[-1])
        return ""

    def review_work(
        self, state: ExecutionState, config: RunnableConfig | None = None
    ) -> ExecutionState:
        """Review whether the completed work adequately addresses the request.

        This node runs after the executor emits an ordinary assistant response without
        requesting any tool calls. It asks the LLM for a structured assessment of the
        work so far. If the work is incomplete, the review rationale is appended as a
        HumanMessage so the executor can continue with explicit feedback. If the work
        is complete, only the structured review state is updated and the graph proceeds
        to recap.
        """
        new_state = deepcopy(state)
        events = self.events(config)
        full_overwrite = False

        new_state, full_overwrite = self.prepare_messages_context(
            new_state, system_prompt=self.executor_prompt
        )

        user_request = self._get_current_user_request(new_state)
        review_message = HumanMessage(content=get_review_prompt(user_request))
        review_messages = list(new_state["messages"]) + [review_message]

        review = invoke_structured(
            self.llm,
            ReviewAssessment,
            review_messages,
            config=self.build_config(tags=["review"]),
            context="execution review assessment",
            fallback=ReviewAssessment(
                is_complete=True,
                reason=(
                    "Review step failed to produce a valid structured "
                    "assessment. Proceeding to recap."
                ),
            ),
            repair=3,
        )

        is_complete = bool(review.is_complete)
        reason = str(review.reason)

        events.emit(
            "Execution Review: Complete"
            if is_complete
            else "Execution Review: Continue",
            stage="review_work",
            approved=is_complete,
            reason=reason,
        )

        if is_complete:
            if full_overwrite:
                return {
                    "messages": Overwrite(new_state["messages"]),
                    "review": review,
                    "symlinkdir": new_state.get("symlinkdir", {}),
                }
            return {"review": review}

        feedback = HumanMessage(
            content=(
                "Review of the current work determined that the current user "
                "request has not yet been adequately addressed. Continue working "
                "and specifically address the following review feedback:\n\n"
                f"{reason}"
            )
        )
        return self.messages_update(
            new_state,
            [feedback],
            full_overwrite=full_overwrite,
            extra={
                "review": review,
                "symlinkdir": new_state.get("symlinkdir", {}),
            },
        )

    def recap(
        self,
        state: ExecutionState,
        config: RunnableConfig | None = None,
    ) -> ExecutionState:
        """Produce a concise summary of the conversation

        This method builds a summarization prompt, invokes the LLM to obtain a compact
        summary of recent interactions and writes debug state when logging is enabled.

        Args:
            state (ExecutionState): The execution state containing message history.

        Returns:
            ExecutionState: A partial update with a single string message containing
                the recap.
        """
        new_state = deepcopy(state)
        events = self.events(config)
        full_overwrite = False

        # 0) Check message history length, summarize, and patch dangling tool calls.
        new_state, full_overwrite = self.prepare_messages_context(
            new_state, system_prompt=self.executor_prompt
        )

        # 1) Construct the summarization message list (system prompt + prior messages).
        recap_message = HumanMessage(content=self.recap_prompt)
        new_state["messages"] = new_state["messages"] + [recap_message]

        # 2) Invoke the LLM to generate a recap; capture content even on failure.
        try:
            response = self.llm.invoke(
                input=new_state["messages"],
                config=self.build_config(tags=["recap"]),
            )
            response_content = response.text
        except Exception:
            try:
                response = self.tool_llm.invoke(
                    input=new_state["messages"],
                    config=self.build_config(tags=["recap"]),
                )
                response_content = response.text
            except Exception as e:
                response_content = f"Response error {e}"
                response = AIMessage(content=response_content)
                events.emit(
                    "Recap failed",
                    stage="recap_error",
                    error_type=type(e).__name__,
                    error=str(e),
                    latest_message=new_state["messages"][-1].text[:2000],
                )

        if full_overwrite:
            # 3) Optionally write state to disk for debugging/auditing.
            new_state["messages"].append(response)
            if self.log_state:
                self.write_state("execution_agent.json", new_state)
            return Overwrite(new_state)
        else:
            if self.log_state:
                new_state["messages"].append(response)
                self.write_state("execution_agent.json", new_state)
            return {"messages": [recap_message, response]}

    def _build_graph(self):
        """Construct and compile the agent's LangGraph state machine."""

        # Keep self.llm unbound for summary/recap calls. The executor loop uses a
        # separate tool-bound model so provider-specific tool transcripts cannot
        # leak into summarization history.
        self.tool_llm = self.tool_llm.bind_tools(self.tools.values())

        # Register nodes:
        # - "agent": LLM planning/execution step
        # - "action": tool dispatch (run_command, write_code, etc.)
        # - "review": structured completeness review before final recap
        # - "recap": summary/finalization step
        self.add_node(self.query_executor, "agent")
        self.add_node(self.tool_node, "action")
        self.add_node(self.review_work, "review")
        self.add_node(self.recap, "recap")

        # Set entrypoint: execution starts with the "agent" node.
        self.graph.set_entry_point("agent")

        # From "agent", either continue with tools or review the completed work,
        # based on presence of tool calls in the last message.
        self.graph.add_conditional_edges(
            "agent",
            self._wrap_cond(should_continue, "should_continue", "execution"),
            {"continue": "action", "review": "review"},
        )

        # After tools run, return control to the agent for the next step.
        self.graph.add_edge("action", "agent")

        # After review, either ask the executor to continue with review feedback or
        # finish with the recap node.
        self.graph.add_conditional_edges(
            "review",
            self._wrap_cond(review_complete, "review_complete", "execution"),
            {"continue": "agent", "recap": "recap"},
        )

        # The graph completes at the "recap" node.
        self.graph.set_finish_point("recap")

    def format_result(self, state: ExecutionState) -> str:
        return state["messages"][-1].text

query_executor(state, runtime, config=None)

Prepare workspace, handle optional symlinks, and invoke the executor LLM.

This method copies the incoming state, ensures a workspace directory exists (creating one with a default name when absent), optionally creates a symlink described by state["symlinkdir"], sets or injects the executor system prompt as the first message, and invokes the bound LLM. When logging is enabled, it persists the pre-invocation state to disk.

Parameters:

Name	Type	Description	Default
state	ExecutionState	The current execution state. Expected keys include: - "messages": Ordered list of System/Human/AI/Tool messages. - "workspace": Optional path to the working directory. - "symlinkdir": Optional dict with "source" and "dest" keys.	required

Returns:

Name	Type	Description
ExecutionState	ExecutionState	Partial state update containing: - "messages": A list with the model's response as the latest entry. - "workspace": The resolved workspace path.

Source code in src/ursa/agents/execution_agent.py

def query_executor(
    self,
    state: ExecutionState,
    runtime: Runtime[AgentContext],
    config: RunnableConfig | None = None,
) -> ExecutionState:
    """Prepare workspace, handle optional symlinks, and invoke the executor LLM.

    This method copies the incoming state, ensures a workspace directory exists
    (creating one with a default name when absent), optionally creates a symlink
    described by state["symlinkdir"], sets or injects the executor system prompt
    as the first message, and invokes the bound LLM. When logging is enabled,
    it persists the pre-invocation state to disk.

    Args:
        state: The current execution state. Expected keys include:
            - "messages": Ordered list of System/Human/AI/Tool messages.
            - "workspace": Optional path to the working directory.
            - "symlinkdir": Optional dict with "source" and "dest" keys.

    Returns:
        ExecutionState: Partial state update containing:
            - "messages": A list with the model's response as the latest entry.
            - "workspace": The resolved workspace path.
    """
    # Add model to the state so it can be passed to tools like the URSA Arxiv or OSTI tools
    new_state = deepcopy(state)
    events = self.events(config)
    new_state.setdefault("symlinkdir", {})

    full_overwrite = False

    # 1.5) Check message history length, summarize, and patch dangling tool calls.
    new_state, full_overwrite = self.prepare_messages_context(
        new_state, system_prompt=self.executor_prompt
    )

    # 2) Optionally create a symlink if symlinkdir is provided and not yet linked.
    sd = new_state.get("symlinkdir")
    if sd and "is_linked" not in sd:
        # symlinkdir structure: {"source": "/path/to/src", "dest": "link/name"}
        symlinkdir = sd

        src = Path(symlinkdir["source"]).expanduser().resolve()
        dst = runtime.context.workspace.joinpath(symlinkdir["dest"])

        # If a file/link already exists at the destination, replace it.
        if dst.exists() or dst.is_symlink():
            dst.unlink()

        # Ensure parent directories for the link exist.
        dst.parent.mkdir(parents=True, exist_ok=True)

        # Create the symlink (tell pathlib if the target is a directory).
        dst.symlink_to(src, target_is_directory=src.is_dir())
        events.emit(
            "Workspace symlink created",
            stage="workspace_symlink",
            source=str(src),
            destination=str(dst),
        )
        new_state["symlinkdir"]["is_linked"] = True
        full_overwrite = True

    # 3) Ensure the executor prompt is the first SystemMessage.
    messages = deepcopy(new_state["messages"])
    if isinstance(messages[0], SystemMessage):
        messages[0] = SystemMessage(content=self.executor_prompt)
    else:
        messages = [SystemMessage(content=self.executor_prompt)] + messages

    # 4) Invoke the LLM with the prepared message sequence.
    try:
        response = self.tool_llm.invoke(
            messages, self.build_config(tags=["agent"])
        )
        new_state["messages"].append(response)
    except Exception as e:  # noqa: BLE001
        response = AIMessage(content=f"Response error {e}")
        msg = new_state["messages"][-1].text
        events.emit(
            "Executor response failed",
            stage="executor_error",
            error_type=type(e).__name__,
            error=str(e),
            latest_message=msg[:2000],
        )
        new_state["messages"].append(response)

    # 5) Optionally persist the pre-invocation state for audit/debugging.
    if self.log_state:
        self.write_state("execution_agent.json", new_state)
    if full_overwrite:
        return {
            "messages": Overwrite(new_state["messages"]),
            "symlinkdir": new_state["symlinkdir"],
        }
    else:
        return {"messages": response, "symlinkdir": new_state["symlinkdir"]}

recap(state, config=None)

Produce a concise summary of the conversation

This method builds a summarization prompt, invokes the LLM to obtain a compact summary of recent interactions and writes debug state when logging is enabled.

Parameters:

Name	Type	Description	Default
state	ExecutionState	The execution state containing message history.	required

Returns:

Name	Type	Description
ExecutionState	ExecutionState	A partial update with a single string message containing the recap.

Source code in src/ursa/agents/execution_agent.py

def recap(
    self,
    state: ExecutionState,
    config: RunnableConfig | None = None,
) -> ExecutionState:
    """Produce a concise summary of the conversation

    This method builds a summarization prompt, invokes the LLM to obtain a compact
    summary of recent interactions and writes debug state when logging is enabled.

    Args:
        state (ExecutionState): The execution state containing message history.

    Returns:
        ExecutionState: A partial update with a single string message containing
            the recap.
    """
    new_state = deepcopy(state)
    events = self.events(config)
    full_overwrite = False

    # 0) Check message history length, summarize, and patch dangling tool calls.
    new_state, full_overwrite = self.prepare_messages_context(
        new_state, system_prompt=self.executor_prompt
    )

    # 1) Construct the summarization message list (system prompt + prior messages).
    recap_message = HumanMessage(content=self.recap_prompt)
    new_state["messages"] = new_state["messages"] + [recap_message]

    # 2) Invoke the LLM to generate a recap; capture content even on failure.
    try:
        response = self.llm.invoke(
            input=new_state["messages"],
            config=self.build_config(tags=["recap"]),
        )
        response_content = response.text
    except Exception:
        try:
            response = self.tool_llm.invoke(
                input=new_state["messages"],
                config=self.build_config(tags=["recap"]),
            )
            response_content = response.text
        except Exception as e:
            response_content = f"Response error {e}"
            response = AIMessage(content=response_content)
            events.emit(
                "Recap failed",
                stage="recap_error",
                error_type=type(e).__name__,
                error=str(e),
                latest_message=new_state["messages"][-1].text[:2000],
            )

    if full_overwrite:
        # 3) Optionally write state to disk for debugging/auditing.
        new_state["messages"].append(response)
        if self.log_state:
            self.write_state("execution_agent.json", new_state)
        return Overwrite(new_state)
    else:
        if self.log_state:
            new_state["messages"].append(response)
            self.write_state("execution_agent.json", new_state)
        return {"messages": [recap_message, response]}

review_work(state, config=None)

Review whether the completed work adequately addresses the request.

This node runs after the executor emits an ordinary assistant response without requesting any tool calls. It asks the LLM for a structured assessment of the work so far. If the work is incomplete, the review rationale is appended as a HumanMessage so the executor can continue with explicit feedback. If the work is complete, only the structured review state is updated and the graph proceeds to recap.

Source code in src/ursa/agents/execution_agent.py

def review_work(
    self, state: ExecutionState, config: RunnableConfig | None = None
) -> ExecutionState:
    """Review whether the completed work adequately addresses the request.

    This node runs after the executor emits an ordinary assistant response without
    requesting any tool calls. It asks the LLM for a structured assessment of the
    work so far. If the work is incomplete, the review rationale is appended as a
    HumanMessage so the executor can continue with explicit feedback. If the work
    is complete, only the structured review state is updated and the graph proceeds
    to recap.
    """
    new_state = deepcopy(state)
    events = self.events(config)
    full_overwrite = False

    new_state, full_overwrite = self.prepare_messages_context(
        new_state, system_prompt=self.executor_prompt
    )

    user_request = self._get_current_user_request(new_state)
    review_message = HumanMessage(content=get_review_prompt(user_request))
    review_messages = list(new_state["messages"]) + [review_message]

    review = invoke_structured(
        self.llm,
        ReviewAssessment,
        review_messages,
        config=self.build_config(tags=["review"]),
        context="execution review assessment",
        fallback=ReviewAssessment(
            is_complete=True,
            reason=(
                "Review step failed to produce a valid structured "
                "assessment. Proceeding to recap."
            ),
        ),
        repair=3,
    )

    is_complete = bool(review.is_complete)
    reason = str(review.reason)

    events.emit(
        "Execution Review: Complete"
        if is_complete
        else "Execution Review: Continue",
        stage="review_work",
        approved=is_complete,
        reason=reason,
    )

    if is_complete:
        if full_overwrite:
            return {
                "messages": Overwrite(new_state["messages"]),
                "review": review,
                "symlinkdir": new_state.get("symlinkdir", {}),
            }
        return {"review": review}

    feedback = HumanMessage(
        content=(
            "Review of the current work determined that the current user "
            "request has not yet been adequately addressed. Continue working "
            "and specifically address the following review feedback:\n\n"
            f"{reason}"
        )
    )
    return self.messages_update(
        new_state,
        [feedback],
        full_overwrite=full_overwrite,
        extra={
            "review": review,
            "symlinkdir": new_state.get("symlinkdir", {}),
        },
    )

ExecutionState

Bases: TypedDict

TypedDict representing the execution agent's mutable run state used by nodes.

Fields: - messages: list of messages (System/Human/AI/Tool). - symlinkdir: optional dict describing a symlink operation (source, dest, is_linked). - review: optional structured review of whether the work is complete.

Source code in src/ursa/agents/execution_agent.py

class ExecutionState(TypedDict):
    """TypedDict representing the execution agent's mutable run state used by nodes.

    Fields:
    - messages: list of messages (System/Human/AI/Tool).
    - symlinkdir: optional dict describing a symlink operation (source, dest,
      is_linked).
    - review: optional structured review of whether the work is complete.
    """

    messages: Annotated[list[AnyMessage], add_messages]
    symlinkdir: dict
    review: NotRequired[ReviewAssessment]
    current_user_request: NotRequired[str]

ReviewAssessment

Bases: BaseModel

Structured assessment of whether execution has addressed the request.

Source code in src/ursa/agents/execution_agent.py

class ReviewAssessment(BaseModel):
    """Structured assessment of whether execution has addressed the request."""

    is_complete: bool = Field(
        description="Whether the work adequately addresses the request."
    )
    reason: str = Field(description="Brief rationale for the review decision.")

command_safe(state)

Return 'safe' if the last command was safe, otherwise 'unsafe'.

Parameters:

Name	Type	Description	Default
state	ExecutionState	The current execution state containing messages and tool calls.	required

Returns: A literal "safe" if no '[UNSAFE]' tags are in the last command, otherwise "unsafe".

Source code in src/ursa/agents/execution_agent.py

def command_safe(state: ExecutionState) -> Literal["safe", "unsafe"]:
    """Return 'safe' if the last command was safe, otherwise 'unsafe'.

    Args:
        state: The current execution state containing messages and tool calls.
    Returns:
        A literal "safe" if no '[UNSAFE]' tags are in the last command,
        otherwise "unsafe".
    """
    index = -1
    message = state["messages"][index]
    # Loop through all the consecutive tool messages in reverse order
    while isinstance(message, ToolMessage):
        if "[UNSAFE]" in message.text:
            return "unsafe"

        index -= 1
        message = state["messages"][index]

    return "safe"

review_complete(state)

Route to recap when review passes, otherwise continue execution.

Source code in src/ursa/agents/execution_agent.py

def review_complete(state: ExecutionState) -> Literal["recap", "continue"]:
    """Route to recap when review passes, otherwise continue execution."""
    review = state.get("review")
    if isinstance(review, ReviewAssessment):
        return "recap" if review.is_complete else "continue"
    if isinstance(review, dict) and review.get("is_complete"):
        return "recap"
    return "continue"

should_continue(state)

Return 'review' if no tool calls in the last message, else 'continue'.

Parameters:

Name	Type	Description	Default
state	ExecutionState	The current execution state containing messages.	required

Returns:

Type	Description
Literal['review', 'continue']	A literal "review" if the last message has no tool calls,
Literal['review', 'continue']	otherwise "continue".

Source code in src/ursa/agents/execution_agent.py

def should_continue(state: ExecutionState) -> Literal["review", "continue"]:
    """Return 'review' if no tool calls in the last message, else 'continue'.

    Args:
        state: The current execution state containing messages.

    Returns:
        A literal "review" if the last message has no tool calls,
        otherwise "continue".
    """
    messages = state["messages"]
    last_message = messages[-1]
    # If there is no tool call, then review the work before finishing.
    if not last_message.tool_calls:
        return "review"
    # Otherwise if there is, we continue.
    else:
        return "continue"

git_agent

Git-aware coding agent with pluggable language support.

GitAgent

Bases: ExecutionAgent

Execution agent with git tools and optional language-specific extensions.

Use directly for language-agnostic git work, or pass language_tools, language_prompt, and safe_codes for a language-specific variant.

Source code in src/ursa/agents/git_agent.py

class GitAgent(ExecutionAgent):
    """Execution agent with git tools and optional language-specific extensions.

    Use directly for language-agnostic git work, or pass ``language_tools``,
    ``language_prompt``, and ``safe_codes`` for a language-specific variant.
    """

    def __init__(
        self,
        llm: BaseChatModel,
        language_tools: list[BaseTool] | None = None,
        language_prompt: str | None = None,
        safe_codes: list[str] | None = None,
        **kwargs,
    ):
        extra_tools: list[BaseTool] = [*GIT_TOOLS, write_code_with_repo]
        if language_tools:
            extra_tools.extend(language_tools)

        super().__init__(
            llm=llm,
            extra_tools=extra_tools,
            safe_codes=safe_codes or [],
            **kwargs,
        )

        self.executor_prompt = compose_git_prompt(language_prompt or "")

        self.remove_tool([
            "run_command",
            "run_web_search",
            "run_osti_search",
            "run_arxiv_search",
        ])

make_git_agent(llm, language=None, language_tools=None, language_prompt=None, safe_codes=None, **kwargs)

Create a GitAgent, optionally with language-specific tools and prompts.

Parameters:

Name	Type	Description	Default
llm	BaseChatModel	The language model to use.	required
language	str | None	Optional language name for registry lookup. If provided and found in LANGUAGE_REGISTRY, its tools/prompt/safe_codes are used as defaults (overridable by explicit parameters). Unknown languages are logged and ignored, defaulting to git-only agent.	None
language_tools	list[BaseTool] | None	Explicit language tools to add. Overrides registry.	None
language_prompt	str | None	Explicit language prompt. Overrides registry.	None
safe_codes	list[str] | None	Explicit safe code list. Overrides registry.	None
**kwargs		Passed to GitAgent constructor.	{}

Returns:

Type	Description
GitAgent	A GitAgent configured with git tools and optionally language-specific
GitAgent	extensions. Works with any file type without requiring explicit
GitAgent	language registration.

Source code in src/ursa/agents/git_agent.py

def make_git_agent(
    llm: BaseChatModel,
    language: str | None = None,
    language_tools: list[BaseTool] | None = None,
    language_prompt: str | None = None,
    safe_codes: list[str] | None = None,
    **kwargs,
) -> GitAgent:
    """Create a GitAgent, optionally with language-specific tools and prompts.

    Args:
        llm: The language model to use.
        language: Optional language name for registry lookup. If provided and
            found in LANGUAGE_REGISTRY, its tools/prompt/safe_codes are used
            as defaults (overridable by explicit parameters). Unknown languages
            are logged and ignored, defaulting to git-only agent.
        language_tools: Explicit language tools to add. Overrides registry.
        language_prompt: Explicit language prompt. Overrides registry.
        safe_codes: Explicit safe code list. Overrides registry.
        **kwargs: Passed to GitAgent constructor.

    Returns:
        A GitAgent configured with git tools and optionally language-specific
        extensions. Works with any file type without requiring explicit
        language registration.
    """
    # Start with explicit parameters (highest priority)
    tools = language_tools
    prompt = language_prompt
    codes = safe_codes

    # Fill in from registry if language is provided and found
    if (
        language
        and language not in (tools or [])
        and language in LANGUAGE_REGISTRY
    ):
        config = LANGUAGE_REGISTRY[language]
        if tools is None:
            tools = config.get("tools")
        if prompt is None:
            prompt = config.get("prompt")
        if codes is None:
            codes = config.get("safe_codes")
    elif language and language not in LANGUAGE_REGISTRY:
        LOGGER.debug(
            "Language %r not in registry; using git-only agent. Available: %s",
            language,
            sorted(LANGUAGE_REGISTRY),
        )

    return GitAgent(
        llm=llm,
        language_tools=tools,
        language_prompt=prompt,
        safe_codes=codes or [],
        **kwargs,
    )

git_go_agent

Git-aware Go coding agent -- backward-compatible wrapper around GitAgent.

GitGoAgent

Bases: GitAgent

Execution agent specialized for git-managed Go repositories.

Tools: - Git: status, diff, log, ls-files, add, commit, switch, create_branch - Go: build, test, vet, mod tidy, linting (golangci-lint with .golangci.yml support) - Code formatting: gofmt

This is a convenience subclass of :class:GitAgent with the Go language tools and prompt pre-configured.

Source code in src/ursa/agents/git_go_agent.py

class GitGoAgent(GitAgent):
    """Execution agent specialized for git-managed Go repositories.

    Tools:
    - Git: status, diff, log, ls-files, add, commit, switch, create_branch
    - Go: build, test, vet, mod tidy, linting (golangci-lint with .golangci.yml support)
    - Code formatting: gofmt

    This is a convenience subclass of :class:`GitAgent` with the Go language
    tools and prompt pre-configured.
    """

    def __init__(self, llm: BaseChatModel, **kwargs):
        super().__init__(
            llm=llm,
            language_tools=GO_TOOLS,
            language_prompt=go_language_prompt,
            safe_codes=["go"],
            **kwargs,
        )

hypothesizer_agent

HypothesizerAgent

Bases: BaseAgent[HypothesizerState]

Maintain a persistent, shareable hypothesis space.

Unlike the legacy/current-review workflow, this agent tracks alternative hypotheses, relative likelihoods, and evidence for/against each hypothesis. The full artifact is stored in den/experiences/<experience_filename> so other agents can bring it back into context with read_experience even if conversational context has been summarized away.

Source code in src/ursa/agents/hypothesizer_agent.py

class HypothesizerAgent(BaseAgent[HypothesizerState]):
    """Maintain a persistent, shareable hypothesis space.

    Unlike the legacy/current-review workflow, this agent tracks alternative
    hypotheses, relative likelihoods, and evidence for/against each hypothesis.
    The full artifact is stored in ``den/experiences/<experience_filename>`` so
    other agents can bring it back into context with ``read_experience`` even if
    conversational context has been summarized away.
    """

    state_type = HypothesizerState

    def __init__(
        self,
        llm: BaseChatModel,
        experience_filename: str = DEFAULT_HYPOTHESIS_EXPERIENCE,
        **kwargs,
    ):
        super().__init__(llm, **kwargs)
        self.experience_filename = self._validate_experience_filename(
            experience_filename
        )

    def _normalize_inputs(self, inputs) -> HypothesizerState:
        if isinstance(inputs, str):
            return HypothesizerState(
                query=inputs,
                new_information=inputs,
                experience_filename=self.experience_filename,
                revision_history=[],
            )

        state = dict(cast(dict[str, Any], inputs))
        if "query" not in state and "question" in state:
            state["query"] = state["question"]
        if "new_information" not in state:
            state["new_information"] = state.get("query", "")
        state.setdefault("experience_filename", self.experience_filename)
        state.setdefault("revision_history", [])
        return cast(HypothesizerState, state)

    def format_query(
        self,
        prompt: str,
        state: HypothesizerState | None = None,
    ) -> HypothesizerState:
        """Treat follow-up prompts as new information for the existing space."""
        if state is None:
            return self._normalize_inputs(prompt)
        updated = dict(state)
        updated["new_information"] = prompt
        updated.setdefault("query", state.get("query", prompt))
        updated.setdefault("experience_filename", self.experience_filename)
        updated.setdefault("revision_history", [])
        return cast(HypothesizerState, updated)

    def format_result(self, result: HypothesizerState) -> str:
        artifact = self._response_text(result.get("hypothesis_space_markdown"))
        if artifact:
            return artifact
        summary = self._response_text(result.get("summary"))
        return (
            summary
            or "Hypothesizer failed to produce a hypothesis-space artifact."
        )

    @staticmethod
    def _validate_experience_filename(filename: str) -> str:
        name = filename.strip()
        path = Path(name)
        if not name:
            raise ValueError("Experience filename must not be empty.")
        if path.is_absolute() or path.name != name or name in {".", ".."}:
            raise ValueError(
                "Experience filename must be a simple relative file name."
            )
        if path.suffix.lower() != ".md":
            raise ValueError("Experience filename must use the .md extension.")
        return name

    @property
    def _experiences_dir(self) -> Path:
        path = self.den / "experiences"
        path.mkdir(parents=True, exist_ok=True)
        return path

    def _experience_path(self, filename: str) -> Path:
        safe_filename = self._validate_experience_filename(filename)
        return self._experiences_dir / safe_filename

    def _read_existing_hypothesis_space(self, filename: str) -> str:
        path = self._experience_path(filename)
        if not path.exists():
            return ""
        try:
            return self._response_text(path.read_text(encoding="utf-8"))
        except OSError:
            return ""

    def _write_hypothesis_space(self, filename: str, content: str) -> Path:
        path = self._experience_path(filename)
        path.write_text(content.rstrip() + "\n", encoding="utf-8")
        return path

    @staticmethod
    def _clean_markdown_text(text: str) -> str:
        """Normalize markdown text returned by model/client layers."""
        text = text.strip()
        if "\\n" in text and text.count("\\n") >= text.count("\n"):
            text = text.replace("\\r\\n", "\n").replace("\\n", "\n")
        return text.strip()

    @classmethod
    def _response_text(cls, value: Any) -> str:
        """Extract user-visible text from LLM response content.

        Some chat models return structured content blocks, e.g. reasoning blocks
        plus text blocks. Avoid ``str(list_of_blocks)`` because that surfaces
        Python/JSON-ish wrapper data to the CLI and hypothesis artifact.
        """
        content = getattr(value, "content", value)
        if content is None:
            return ""
        if isinstance(content, str):
            raw_text = content.strip()
            if (
                raw_text.startswith(("[", "{"))
                and "'type'" in raw_text
                and "'text'" in raw_text
            ):
                try:
                    parsed = ast.literal_eval(raw_text)
                except (SyntaxError, ValueError):
                    return cls._clean_markdown_text(raw_text)
                parsed_text = cls._response_text(parsed)
                return parsed_text or cls._clean_markdown_text(raw_text)
            if (
                raw_text.startswith(("[", "{"))
                and '"type"' in raw_text
                and '"text"' in raw_text
            ):
                try:
                    parsed = ast.literal_eval(raw_text)
                except (SyntaxError, ValueError):
                    return cls._clean_markdown_text(raw_text)
                parsed_text = cls._response_text(parsed)
                return parsed_text or cls._clean_markdown_text(raw_text)
            return cls._clean_markdown_text(raw_text)
        if isinstance(content, list):
            parts = [cls._response_text(item) for item in content]
            return "\n\n".join(part for part in parts if part).strip()
        if isinstance(content, dict):
            block_type = content.get("type")
            if block_type in {"text", "output_text"} and "text" in content:
                return cls._response_text(content["text"])
            if block_type == "reasoning":
                return ""
            if "text" in content:
                return cls._response_text(content["text"])
            if "content" in content:
                return cls._response_text(content["content"])
            return ""
        return cls._clean_markdown_text(str(content))

    def _fallback_hypothesis_space(
        self,
        *,
        query: str,
        new_information: str,
        context: str,
        previous: str,
        now: str,
    ) -> str:
        basis = new_information or query or "the current user question"
        context_note = (
            context or "No additional cross-agent context was provided."
        )
        previous_note = (
            "A prior hypothesis-space artifact existed and should remain part "
            "of the ongoing context."
            if previous.strip()
            else "No previous hypothesis-space artifact was found."
        )
        return f"""# Hypothesis Space

## Question / Topic

{query or basis}

## Current Update

{basis}

## Additional Context

{context_note}

## Hypotheses

### H1: Primary working hypothesis

- **Relative likelihood:** 0.34
- **Rationale:** This is the most direct explanation currently available from the supplied prompt/context.
- **Evidence for:**
  - The current user-provided information is consistent with this hypothesis.
- **Evidence against:**
  - No explicit contradictory evidence has been recorded yet.
- **Assumptions / uncertainties:**
  - This likelihood is provisional and should be updated as more evidence is gathered.

### H2: Alternative mechanism or explanation

- **Relative likelihood:** 0.33
- **Rationale:** A plausible alternative may explain the same observations through a different causal path.
- **Evidence for:**
  - The current evidence does not rule it out.
- **Evidence against:**
  - No specific supporting evidence has been isolated yet.
- **Assumptions / uncertainties:**
  - Additional targeted evidence is needed to compare it against H1.

### H3: Null, mixed, or confounded explanation

- **Relative likelihood:** 0.33
- **Rationale:** The available observations may be incomplete, confounded, or explained by multiple factors.
- **Evidence for:**
  - {previous_note}
- **Evidence against:**
  - No decisive evidence yet.
- **Assumptions / uncertainties:**
  - More precise observations, source documents, experiments, or logs would reduce uncertainty.

## Evidence Log

- **E1:** {basis}
  - Supports: H1, H2, H3 to varying degrees.
  - Contradicts: none recorded yet.
  - Strength: provisional.

## Change Summary

Initialized or updated the hypothesis space with the latest user-provided information. Treat all likelihoods as provisional until additional evidence is gathered.

## Recommended Next Evidence

- Identify observations that would distinguish H1 from H2.
- Look for evidence that would falsify the primary working hypothesis.
- Record source documents, commands, results, or experiments as future evidence updates.
"""

    def _build_update_prompt(
        self,
        *,
        query: str,
        new_information: str,
        context: str,
        previous: str,
        now: str,
        filename: str,
    ) -> list:
        system = SystemMessage(
            content=(
                "You are a persistent hypothesis-space maintainer for URSA. "
                "Your job is to maintain a structured set of competing hypotheses, "
                "relative likelihoods, and evidence for/against each hypothesis. "
                "Return only Markdown. Do not use markdown code fences around the whole response. "
                "The artifact will be written as an experience file so other agents can read it later."
            )
        )
        human = HumanMessage(
            content=f"""Update the hypothesis-space artifact.

The artifact will be saved to the configured hypothesis-space experience file.
Do not include implementation metadata such as file paths, timestamps, JSON, or HTML comments in the output.

Question / topic:
{query}

New user information / evidence / instruction:
{new_information}

Additional context from another agent behavior or user-provided notes:
{context or "(none provided)"}

Previous hypothesis-space artifact:
{previous or "(none found; initialize a new hypothesis space)"}

Requirements:
- Maintain clear hypothesis IDs such as H1, H2, H3.
- Keep or update relative likelihoods.
- State whether likelihoods are mutually exclusive probabilities or independent plausibility scores.
- Track evidence for and against each hypothesis.
- Preserve useful prior evidence unless contradicted.
- Explain exactly what changed in this update.
- Include recommended next evidence or work that chat/execution agents could gather.
- Keep the artifact concise enough to be read back into context later, but complete enough to be useful.
"""
        )
        return [system, human]

    def _summarize_artifact(self, artifact: str, filename: str) -> str:
        lines = [line.strip() for line in artifact.splitlines() if line.strip()]
        hypotheses = [line for line in lines if line.startswith("### H")]
        if hypotheses:
            hyp_text = "; ".join(hypotheses[:5])
            return (
                f"Updated hypothesis space in experiences/{filename}. "
                f"Current hypotheses: {hyp_text}"
            )
        return f"Updated hypothesis space in experiences/{filename}."

    def update_hypothesis_space(
        self,
        state: HypothesizerState,
        config: RunnableConfig | None = None,
    ) -> HypothesizerState:
        events = self.events(config)
        filename = self._validate_experience_filename(
            state.get("experience_filename", self.experience_filename)
        )
        query = state.get("query") or state.get("new_information", "")
        new_information = state.get("new_information") or query
        context = state.get("context", "")
        now = datetime.now(UTC).isoformat()
        previous = self._read_existing_hypothesis_space(filename)

        events.emit(
            "Updating hypothesis space",
            stage="hypothesis_update",
            experience_filename=filename,
            has_existing_artifact=bool(previous.strip()),
        )

        messages = self._build_update_prompt(
            query=query,
            new_information=new_information,
            context=context,
            previous=previous,
            now=now,
            filename=filename,
        )
        response = self.llm.invoke(messages, config=config)
        artifact = self._response_text(response)

        if not artifact or artifact.lower() in {"ok", "stub"}:
            artifact = self._fallback_hypothesis_space(
                query=query,
                new_information=new_information,
                context=context,
                previous=previous,
                now=now,
            )

        if not artifact.lstrip().startswith("#"):
            artifact = f"# Hypothesis Space\n\n{artifact}"

        artifact = artifact.rstrip()

        path = self._write_hypothesis_space(filename, artifact)
        summary = self._summarize_artifact(artifact, filename)
        revision_entry = (
            f"{now}: updated {filename} with new information: "
            f"{new_information[:160]}"
        )
        revision_history = list(state.get("revision_history", []))
        revision_history.append(revision_entry)

        events.emit(
            "Hypothesis space updated",
            stage="hypothesis_update",
            experience_filename=filename,
            output_path=str(path),
            summary=summary,
        )

        return HypothesizerState(
            query=query,
            new_information=new_information,
            context=context,
            experience_filename=filename,
            hypothesis_space_markdown=artifact,
            summary=summary,
            revision_history=revision_history,
            last_updated=now,
        )

    def _build_graph(self):
        self.add_node(self.update_hypothesis_space, "update_hypothesis_space")
        self.graph.set_entry_point("update_hypothesis_space")
        self.graph.set_finish_point("update_hypothesis_space")

format_query(prompt, state=None)

Treat follow-up prompts as new information for the existing space.

Source code in src/ursa/agents/hypothesizer_agent.py

def format_query(
    self,
    prompt: str,
    state: HypothesizerState | None = None,
) -> HypothesizerState:
    """Treat follow-up prompts as new information for the existing space."""
    if state is None:
        return self._normalize_inputs(prompt)
    updated = dict(state)
    updated["new_information"] = prompt
    updated.setdefault("query", state.get("query", prompt))
    updated.setdefault("experience_filename", self.experience_filename)
    updated.setdefault("revision_history", [])
    return cast(HypothesizerState, updated)

HypothesizerState

Bases: TypedDict

State for persistent hypothesis-space maintenance.

The durable hypothesis space is intentionally offloaded to an experience file rather than kept entirely in graph state. This makes it easy for other URSA agent behaviors, such as chat and execution, to read the current hypothesis space back into context via the existing experience tools.

Source code in src/ursa/agents/hypothesizer_agent.py

class HypothesizerState(TypedDict, total=False):
    """State for persistent hypothesis-space maintenance.

    The durable hypothesis space is intentionally offloaded to an experience file
    rather than kept entirely in graph state. This makes it easy for other URSA
    agent behaviors, such as chat and execution, to read the current hypothesis
    space back into context via the existing experience tools.
    """

    query: str
    """Original or current user question/topic."""

    new_information: str
    """New evidence, clarification, or instruction to incorporate."""

    context: str
    """Optional additional context from another agent behavior or user-provided notes."""

    experience_filename: str
    """Markdown experience file used as the durable hypothesis artifact."""

    hypothesis_space_markdown: str
    """Latest hypothesis-space artifact markdown."""

    summary: str
    """Compact state summary of the current hypothesis space."""

    revision_history: list[str]
    """Short descriptions of updates made during this run/thread."""

    last_updated: str
    """ISO timestamp for the latest update."""

context instance-attribute

Optional additional context from another agent behavior or user-provided notes.

experience_filename instance-attribute

Markdown experience file used as the durable hypothesis artifact.

hypothesis_space_markdown instance-attribute

Latest hypothesis-space artifact markdown.

last_updated instance-attribute

ISO timestamp for the latest update.

new_information instance-attribute

New evidence, clarification, or instruction to incorporate.

query instance-attribute

Original or current user question/topic.

revision_history instance-attribute

Short descriptions of updates made during this run/thread.

summary instance-attribute

Compact state summary of the current hypothesis space.

should_continue(state)

Compatibility helper for callers that imported the old symbol.

Source code in src/ursa/agents/hypothesizer_agent.py

def should_continue(state: HypothesizerState) -> Literal["finish"]:
    """Compatibility helper for callers that imported the old symbol."""
    return "finish"

lammps_agent

LammpsAgent

Bases: BaseAgent[LammpsState]

Source code in src/ursa/agents/lammps_agent.py

class LammpsAgent(BaseAgent[LammpsState]):
    state_type = LammpsState

    def __init__(
        self,
        llm: BaseChatModel,
        potential_files: Optional[list[str]] = None,
        pair_style: Optional[str] = None,
        pair_coeff: Optional[str] = None,
        max_potentials: int = 5,
        max_fix_attempts: int = 10,
        find_potential_only: bool = False,
        data_file: str = None,
        data_max_lines: int = 50,
        ngpus: int = -1,
        mpi_procs: int = 8,
        workspace: str = "./workspace",
        lammps_cmd: str = "lmp_mpi",
        mpirun_cmd: str = "mpirun",
        tiktoken_model: str = "gpt-5-mini",
        max_tokens: int = 200000,
        summarize_results: bool = True,
        **kwargs,
    ):
        if not working:
            raise ImportError(
                "LAMMPS agent requires the atomman and trafilatura dependencies. These can be installed using 'pip install ursa-ai[lammps]' or, if working from a local installation, 'pip install -e .[lammps]' ."
            )

        super().__init__(llm, **kwargs)

        self.user_potential_files = potential_files
        self.user_pair_style = pair_style
        self.user_pair_coeff = pair_coeff
        self.use_user_potential = (
            potential_files is not None
            and pair_style is not None
            and pair_coeff is not None
        )

        self.max_potentials = max_potentials
        self.max_fix_attempts = max_fix_attempts
        self.find_potential_only = find_potential_only
        self.data_file = data_file
        self.data_max_lines = data_max_lines
        self.ngpus = ngpus
        self.mpi_procs = mpi_procs
        self.lammps_cmd = lammps_cmd
        self.mpirun_cmd = mpirun_cmd
        self.tiktoken_model = tiktoken_model
        self.max_tokens = max_tokens
        self.summarize_results = summarize_results

        self.pair_styles = [
            "eam",
            "eam/alloy",
            "eam/fs",
            "meam",
            "adp",
            "kim",
            "snap",
            "quip",
            "mlip",
            "pace",
            "nep",
        ]

        self.workspace = workspace
        os.makedirs(self.workspace, exist_ok=True)

        self.str_parser = StrOutputParser()

        self.summ_chain = (
            ChatPromptTemplate.from_template(
                "Here is some data about an interatomic potential: {metadata}\n\n"
                "Briefly summarize why it could be useful for this task: {simulation_task}."
            )
            | self.llm
            | self.str_parser
        )

        self.choose_chain = (
            ChatPromptTemplate.from_template(
                "Here are the summaries of a certain number of interatomic potentials: {summaries_combined}\n\n"
                "Pick one potential which would be most useful for this task: {simulation_task}.\n\n"
                "Return your answer **only** as valid JSON, with no extra text or formatting.\n\n"
                "Use this exact schema:\n"
                "{{\n"
                '  "Chosen index": <int>,\n'
                '  "rationale": "<string>",\n'
                '  "Potential name": "<string>"\n'
                "}}\n"
            )
            | self.llm
            | self.str_parser
        )

        self.author_chain = (
            ChatPromptTemplate.from_template(
                "Your task is to write a LAMMPS input file for this purpose: {simulation_task}.\n"
                "Note that all potential files are in the './' directory.\n"
                "Here is some information about the pair_style and pair_coeff that might be useful in writing the input file: {pair_info}.\n"
                "If a template for the input file is provided, you should adapt it appropriately to meet the task requirements.\n"
                "Template provided (if any): {template}\n"
                "If a data file is provided, use it in the input script via the 'read_data' command.\n"
                "Name of data file (if any): {data_file}\n"
                "First few lines of data file (if any):\n{data_content}\n"
                "Ensure that all logs are recorded in a './log.lammps' file.\n"
                "To create the log file, use may use the 'log ./log.lammps' command. \n"
                "Return your answer **only** as valid JSON, with no extra text or formatting.\n"
                "IMPORTANT: Properly escape all special characters in the input_script string (use \\n for newlines, \\\\ for backslashes, etc.).\n"
                "Use this exact schema:\n"
                "{{\n"
                '  "input_script": "<string>"\n'
                "}}\n"
            )
            | self.llm
            | self.str_parser
        )

        self.fix_chain = (
            ChatPromptTemplate.from_template(
                "You are part of a larger scientific workflow whose purpose is to accomplish this task: {simulation_task}\n"
                "Multiple attempts at writing and running a LAMMPS input file have been made.\n"
                "Here is the run history across attempts (each includes the input script and its stdout/stderr):{err_message}\n"
                "Use the history to identify what changed between attempts and avoid repeating failed approaches.\n"
                "Your task is to write a new input file that resolves the latest error.\n"
                "Note that all potential files are in the './' directory.\n"
                "Here is some information about the pair_style and pair_coeff that might be useful in writing the input file: {pair_info}.\n"
                "If a template for the input file is provided, you should adapt it appropriately to meet the task requirements.\n"
                "Template provided (if any): {template}\n"
                "If a data file is provided, use it in the input script via the 'read_data' command.\n"
                "Name of data file (if any): {data_file}\n"
                "First few lines of data file (if any):\n{data_content}\n"
                "Ensure that all logs are recorded in a './log.lammps' file.\n"
                "To create the log file, use may use the 'log ./log.lammps' command. \n"
                "Return your answer **only** as valid JSON, with no extra text or formatting.\n"
                "IMPORTANT: Properly escape all special characters in the input_script string (use \\n for newlines, \\\\ for backslashes, etc.).\n"
                "Use this exact schema:\n"
                "{{\n"
                '  "input_script": "<string>"\n'
                "}}\n"
            )
            | self.llm
            | self.str_parser
        )

    def _section(self, title: str):
        LOGGER.info("%s", title)

    def _panel(self, title: str, body: str, style: str = "cyan"):
        LOGGER.info("%s\n%s", title, body)

    def _code_panel(
        self,
        title: str,
        code: str,
        language: str = "bash",
        style: str = "magenta",
    ):
        LOGGER.info("%s\n%s", title, code)

    def _diff_panel(self, old: str, new: str, title: str = "LAMMPS input diff"):
        diff = "\n".join(
            difflib.unified_diff(
                old.splitlines(),
                new.splitlines(),
                fromfile="in.lammps (before)",
                tofile="in.lammps (after)",
                lineterm="",
            )
        )
        if not diff.strip():
            diff = "(no changes)"
        LOGGER.info("%s\n%s", title, diff)

    @staticmethod
    def _safe_json_loads(s: str) -> dict[str, Any]:
        s = s.strip()
        if s.startswith("```"):
            s = s.strip("`")
            i = s.find("\n")
            if i != -1:
                s = s[i + 1 :].strip()
        return json.loads(s)

    def _read_and_trim_data_file(
        self,
        data_file_path: str,
        config: RunnableConfig | None = None,
    ) -> str:
        """Read LAMMPS data file and trim to token limit for LLM context."""
        events = self.events(config)
        if os.path.exists(data_file_path):
            with open(data_file_path, "r") as f:
                content = f.read()
            lines = content.splitlines()
            if len(lines) > self.data_max_lines:
                content = "\n".join(lines[: self.data_max_lines])
                events.emit(
                    "Data file trimmed",
                    stage="read_data",
                    path=data_file_path,
                    original_lines=len(lines),
                    retained_lines=self.data_max_lines,
                )
            return content
        else:
            return "Could not read data file."

    def _copy_data_file(
        self,
        data_file_path: str,
        config: RunnableConfig | None = None,
    ) -> str:
        """Copy data file to workspace and return new path."""
        if not os.path.exists(data_file_path):
            raise FileNotFoundError(f"Data file not found: {data_file_path}")

        filename = os.path.basename(data_file_path)
        dest_path = os.path.join(self.workspace, filename)
        os.system(f"cp {data_file_path} {dest_path}")
        self.events(config).emit(
            "Data file copied to workspace",
            stage="copy_data",
            path=dest_path,
        )
        return dest_path

    def _copy_user_potential_files(
        self,
        config: RunnableConfig | None = None,
    ):
        """Copy user-provided potential files to workspace."""
        events = self.events(config)
        events.emit(
            "Copying user-provided potential files",
            stage="copy_potentials",
        )
        for pot_file in self.user_potential_files:
            if not os.path.exists(pot_file):
                raise FileNotFoundError(f"Potential file not found: {pot_file}")

            filename = os.path.basename(pot_file)
            dest_path = os.path.join(self.workspace, filename)

            try:
                os.system(f"cp {pot_file} {dest_path}")
                events.emit(
                    "Potential file copied to workspace",
                    stage="copy_potentials",
                    path=dest_path,
                )
            except Exception:
                LOGGER.exception("Error copying %s", filename)
                raise

    def _create_user_potential_wrapper(
        self,
        state: LammpsState,
        config: RunnableConfig | None = None,
    ) -> LammpsState:
        """Create a wrapper object for user-provided potential to match atomman interface."""
        self._copy_user_potential_files(config)

        # Create a simple object that mimics the atomman potential interface
        class UserPotential:
            def __init__(self, pair_style, pair_coeff):
                self._pair_style = pair_style
                self._pair_coeff = pair_coeff

            def pair_info(self):
                return f"pair_style {self._pair_style}\npair_coeff {self._pair_coeff}"

        user_potential = UserPotential(
            self.user_pair_style, self.user_pair_coeff
        )

        return {
            **state,
            "chosen_potential": user_potential,
            "fix_attempts": 0,
            "run_history": [],
        }

    def _fetch_and_trim_text(self, url: str) -> str:
        downloaded = trafilatura.fetch_url(url)
        if not downloaded:
            return "No metadata available"
        text = trafilatura.extract(
            downloaded,
            include_comments=False,
            include_tables=True,
            include_links=False,
            favor_recall=True,
        )
        if not text:
            return "No metadata available"
        text = text.strip()
        try:
            enc = tiktoken.encoding_for_model(self.tiktoken_model)
            toks = enc.encode(text)
            if len(toks) > self.max_tokens:
                toks = toks[: self.max_tokens]
                text = enc.decode(toks)
        except Exception:
            pass
        return text

    def _entry_router(
        self,
        state: LammpsState,
        config: RunnableConfig | None = None,
    ) -> dict:
        events = self.events(config)
        # Check if using user-provided potential
        if self.use_user_potential:
            if self.find_potential_only:
                raise Exception(
                    "Cannot set find_potential_only=True when providing your own potential!"
                )
            events.emit(
                "Using user-provided potential files",
                stage="entry",
            )

        if self.find_potential_only and state.get("chosen_potential"):
            raise Exception(
                "You cannot set find_potential_only=True and also specify your own potential!"
            )

        if self.data_file:
            try:
                self._copy_data_file(self.data_file, config)
            except Exception as e:
                LOGGER.warning("Could not process data file: %s", e)

        if not state.get("chosen_potential"):
            self.potential_summaries_dir = os.path.join(
                self.workspace, "potential_summaries"
            )
            os.makedirs(self.potential_summaries_dir, exist_ok=True)
        return {}

    def _find_potentials(self, state: LammpsState) -> LammpsState:
        db = am.library.Database(remote=True)
        matches = db.get_lammps_potentials(
            pair_style=self.pair_styles, elements=state["elements"]
        )

        return {
            **state,
            "matches": list(matches),
            "idx": 0,
            "summaries": [],
            "full_texts": [],
            "fix_attempts": 0,
            "run_history": [],
        }

    def _should_summarize(
        self,
        state: LammpsState,
        config: RunnableConfig | None = None,
    ) -> str:
        matches = state.get("matches", [])
        i = state.get("idx", 0)
        if not matches:
            self.events(config).emit(
                "No potentials found in NIST for this task",
                stage="find_potentials",
            )
            return "done_no_matches"
        if i < min(self.max_potentials, len(matches)):
            return "summarize_one"
        return "summarize_done"

    def _summarize_one(self, state: LammpsState) -> LammpsState:
        i = state["idx"]
        self._section(f"Summarizing potential #{i}")
        match = state["matches"][i]
        md = match.metadata()

        if md.get("comments") is None:
            text = "No metadata available"
            summary = "No summary available"
        else:
            lines = md["comments"].split("\n")
            url = lines[1] if len(lines) > 1 else ""
            text = (
                self._fetch_and_trim_text(url)
                if url
                else "No metadata available"
            )
            summary = self.summ_chain.invoke({
                "metadata": text,
                "simulation_task": state["simulation_task"],
            })

        summary_file = os.path.join(
            self.potential_summaries_dir, "potential_" + str(i) + ".txt"
        )
        with open(summary_file, "w") as f:
            f.write(summary)

        return {
            **state,
            "idx": i + 1,
            "summaries": [*state["summaries"], summary],
            "full_texts": [*state["full_texts"], text],
        }

    def _build_summaries(self, state: LammpsState) -> LammpsState:
        parts = []
        for i, s in enumerate(state["summaries"]):
            rec = state["matches"][i]
            parts.append(f"\nSummary of potential #{i}: {rec.id}\n{s}\n")
        return {**state, "summaries_combined": "".join(parts)}

    def _choose(self, state: LammpsState) -> LammpsState:
        self._section("Choosing potential")
        choice = self.choose_chain.invoke({
            "summaries_combined": state["summaries_combined"],
            "simulation_task": state["simulation_task"],
        })
        choice_dict = self._safe_json_loads(choice)
        chosen_index = int(choice_dict["Chosen index"])

        chosen_potential = state["matches"][chosen_index]

        self._panel(
            "Chosen Potential",
            f"[bold]Index:[/bold] {chosen_index}\n[bold]ID:[/bold] {chosen_potential.id}\n\n[bold]Rationale:[/bold]\n{choice_dict['rationale']}",
            style="green",
        )

        out_file = os.path.join(self.potential_summaries_dir, "Rationale.txt")
        with open(out_file, "w") as f:
            f.write(f"Chosen potential #{chosen_index}")
            f.write("\n")
            f.write("Rationale for choosing this potential:")
            f.write("\n")
            f.write(choice_dict["rationale"])

        return {**state, "chosen_potential": chosen_potential}

    def _route_after_summarization(self, state: LammpsState) -> str:
        if self.find_potential_only:
            return "Exit"
        return "continue_author"

    def _author(
        self,
        state: LammpsState,
        config: RunnableConfig | None = None,
    ) -> LammpsState:
        self._section("First attempt at writing LAMMPS input file")

        if not self.use_user_potential:
            state["chosen_potential"].download_files(self.workspace)
        pair_info = state["chosen_potential"].pair_info()

        data_content = ""
        if self.data_file:
            data_content = self._read_and_trim_data_file(
                self.data_file,
                config,
            )

        authored_json = self.author_chain.invoke({
            "simulation_task": state["simulation_task"],
            "pair_info": pair_info,
            "template": state["template"],
            "data_file": self.data_file,
            "data_content": data_content,
        })
        script_dict = self._safe_json_loads(authored_json)
        input_script = script_dict["input_script"]
        with open(os.path.join(self.workspace, "in.lammps"), "w") as f:
            f.write(input_script)

        self._section("Authored LAMMPS input")
        self._code_panel(
            "in.lammps", input_script, language="bash", style="magenta"
        )

        return {**state, "input_script": input_script}

    def _run_lammps(
        self,
        state: LammpsState,
        config: RunnableConfig | None = None,
    ) -> LammpsState:
        self._section("Running LAMMPS")

        if self.ngpus >= 0:
            result = subprocess.run(
                [
                    self.mpirun_cmd,
                    "-np",
                    str(self.mpi_procs),
                    self.lammps_cmd,
                    "-in",
                    "in.lammps",
                    "-k",
                    "on",
                    "g",
                    str(self.ngpus),
                    "-sf",
                    "kk",
                    "-pk",
                    "kokkos",
                    "neigh",
                    "half",
                    "newton",
                    "on",
                ],
                cwd=self.workspace,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=False,
            )
            self.events(config).emit(
                "LAMMPS command finished",
                stage="run",
                returncode=result.returncode,
                stdout_chars=len(result.stdout or ""),
                stderr_chars=len(result.stderr or ""),
            )
        else:
            result = subprocess.run(
                [
                    self.mpirun_cmd,
                    "-np",
                    str(self.mpi_procs),
                    self.lammps_cmd,
                    "-in",
                    "in.lammps",
                ],
                cwd=self.workspace,
                stdout=subprocess.PIPE,
                stderr=subprocess.PIPE,
                text=True,
                check=False,
            )

        status_style = "green" if result.returncode == 0 else "red"
        self._panel(
            "Run Result",
            f"returncode = {result.returncode}",
            style=status_style,
        )

        if result.returncode != 0:
            err_view = (
                result.stderr.strip() + "\n" + result.stdout.strip()
            ).strip() or "(no output captured)"
            self._panel("Run error/output", err_view[-6000:], style="red")

        hist = list(state.get("run_history", []))
        hist.append({
            "attempt": state.get("fix_attempts", 0),
            "input_script": state.get("input_script", ""),
            "returncode": result.returncode,
            "stdout": result.stdout,
            "stderr": result.stderr,
        })

        return {
            **state,
            "run_returncode": result.returncode,
            "run_stdout": result.stdout,
            "run_stderr": result.stderr,
            "run_history": hist,
        }

    def _route_run(self, state: LammpsState) -> str:
        rc = state.get("run_returncode", 0)
        attempts = state.get("fix_attempts", 0)
        if rc == 0:
            self._section("LAMMPS run successful! Exiting...")
            return "done_success"
        if attempts < self.max_fix_attempts:
            self._section(
                "LAMMPS run Failed. Attempting to rewrite input file..."
            )
            return "need_fix"
        self._section(
            "LAMMPS run Failed and maximum fix attempts reached. Exiting.."
        )
        return "done_failed"

    def _fix(
        self,
        state: LammpsState,
        config: RunnableConfig | None = None,
    ) -> LammpsState:
        pair_info = state["chosen_potential"].pair_info()

        hist = state.get("run_history", [])
        if not hist:
            hist = [
                {
                    "attempt": state.get("fix_attempts", 0),
                    "input_script": state.get("input_script", ""),
                    "returncode": state.get("run_returncode"),
                    "stdout": state.get("run_stdout", ""),
                    "stderr": state.get("run_stderr", ""),
                }
            ]

        parts = []
        for h in hist:
            parts.append(
                "=== Attempt {attempt} | returncode={returncode} ===\n"
                "--- input_script ---\n{input_script}\n"
                "--- stdout ---\n{stdout}\n"
                "--- stderr ---\n{stderr}\n".format(**h)
            )
        err_blob = "\n".join(parts)

        data_content = ""
        if self.data_file:
            data_content = self._read_and_trim_data_file(
                self.data_file,
                config,
            )

        fixed_json = self.fix_chain.invoke({
            "simulation_task": state["simulation_task"],
            "err_message": err_blob,
            "pair_info": pair_info,
            "template": state["template"],
            "data_file": self.data_file,
            "data_content": data_content,
        })
        script_dict = self._safe_json_loads(fixed_json)

        new_input = script_dict["input_script"]
        old_input = state["input_script"]
        self._diff_panel(old_input, new_input)

        with open(os.path.join(self.workspace, "in.lammps"), "w") as f:
            f.write(new_input)

        return {
            **state,
            "input_script": new_input,
            "fix_attempts": state.get("fix_attempts", 0) + 1,
        }

    def _summarize(
        self,
        state: LammpsState,
        config: RunnableConfig | None = None,
    ) -> LammpsState:
        self._section(
            "Now handing things off to execution agent for summarization/visualization"
        )

        executor = ExecutionAgent(llm=self.llm)

        exe_plan = f"""
        You are part of a larger scientific workflow whose purpose is to accomplish this task: {state["simulation_task"]}
        A LAMMPS simulation has been done and the output is located in the file 'log.lammps'.
        Summarize the contents of this file in a markdown document. Include a plot, if relevent.
        """

        exe_results = executor.invoke({
            "messages": [HumanMessage(content=exe_plan)],
            "workspace": self.workspace,
        })

        for x in exe_results["messages"]:
            self.events(config).emit(
                "Execution summary message received",
                stage="summarize_results",
                result_chars=len(str(x.content)),
            )

        return state

    def _post_run(self, state: LammpsState) -> LammpsState:
        return state

    def _build_graph(self):
        self.add_node(self._entry_router)
        self.add_node(self._find_potentials)
        self.add_node(self._summarize_one)
        self.add_node(self._build_summaries)
        self.add_node(self._choose)
        self.add_node(self._create_user_potential_wrapper)
        self.add_node(self._author)
        self.add_node(self._run_lammps)
        self.add_node(self._fix)
        self.add_node(self._post_run)
        self.add_node(self._summarize)

        self.graph.set_entry_point("_entry_router")

        self.graph.add_conditional_edges(
            "_entry_router",
            lambda state: "user_potential"
            if self.use_user_potential
            else (
                "user_choice"
                if state.get("chosen_potential")
                else "agent_choice"
            ),
            {
                "user_potential": "_create_user_potential_wrapper",
                "user_choice": "_author",
                "agent_choice": "_find_potentials",
            },
        )

        self.graph.add_conditional_edges(
            "_find_potentials",
            self._should_summarize,
            {
                "summarize_one": "_summarize_one",
                "summarize_done": "_build_summaries",
                "done_no_matches": END,
            },
        )

        self.graph.add_conditional_edges(
            "_summarize_one",
            self._should_summarize,
            {
                "summarize_one": "_summarize_one",
                "summarize_done": "_build_summaries",
            },
        )

        self.graph.add_edge("_build_summaries", "_choose")

        self.graph.add_conditional_edges(
            "_choose",
            self._route_after_summarization,
            {
                "continue_author": "_author",
                "Exit": END,
            },
        )

        self.graph.add_edge("_create_user_potential_wrapper", "_author")
        self.graph.add_edge("_author", "_run_lammps")

        self.graph.add_conditional_edges(
            "_run_lammps",
            self._route_run,
            {
                "need_fix": "_fix",
                "done_success": "_post_run",
                "done_failed": END,
            },
        )

        self.graph.add_edge("_fix", "_run_lammps")

        self.graph.add_conditional_edges(
            "_post_run",
            lambda _: "summarize" if self.summarize_results else "skip",
            {
                "summarize": "_summarize",
                "skip": END,
            },
        )

        self.graph.add_edge("_summarize", END)

mp_agent

MaterialsProjectAgent

Bases: BaseAgent

Source code in src/ursa/agents/mp_agent.py

class MaterialsProjectAgent(BaseAgent):
    def __init__(
        self,
        llm: BaseChatModel,
        summarize: bool = True,
        max_results: int = 3,
        database_path: str = "mp_database",
        summaries_path: str = "mp_summaries",
        **kwargs,
    ):
        if not working:
            raise ImportError(
                "MP Agent requires special dependencies. These can be installed using 'pip install ursa-ai[mp]' or, if working from a local installation, 'pip install -e .[mp]' ."
            )

        super().__init__(llm, **kwargs)
        self.summarize = summarize
        self.max_results = max_results
        self.database_path = self.den.joinpath(database_path)
        self.summaries_path = self.den.joinpath(summaries_path)

        self.database_path.mkdir(parents=True, exist_ok=True)
        self.summaries_path.mkdir(parents=True, exist_ok=True)

    def _fetch_node(self, state: dict) -> dict:
        f = state["query"]
        els = f["elements"]  # e.g. ["Ga","In"]
        bg = (f["band_gap_min"], f["band_gap_max"])
        e_above_hull = (0, 0)  # only on-hull (stable)
        mats = []
        with MPRester() as mpr:
            # get ALL matching materials…
            all_results = mpr.materials.summary.search(
                elements=els,
                band_gap=bg,
                energy_above_hull=e_above_hull,
                is_stable=True,  # equivalent filter
            )
            # …then take only the first `max_results`
            for doc in all_results[: self.max_results]:
                mid = doc.material_id
                data = doc.dict()
                # cache to disk
                path = os.path.join(self.database_path, f"{mid}.json")
                if not os.path.exists(path):
                    with open(path, "w") as f:
                        json.dump(data, f, indent=2)
                mats.append({"material_id": mid, "metadata": data})

        return {**state, "materials": mats}

    def _summarize_node(self, state: dict) -> dict:
        """Summarize each material via LLM over its metadata."""
        # prompt template
        prompt = ChatPromptTemplate.from_template("""
You are a materials-science assistant. Given the following metadata about a material, produce a concise summary focusing on its key properties:

{metadata}
        """)
        chain = prompt | self.llm | StrOutputParser()

        summaries = [None] * len(state["materials"])

        def process(i, mat):
            mid = mat["material_id"]
            meta = mat["metadata"]
            # flatten metadata to text
            text = "\n".join(f"{k}: {v}" for k, v in meta.items())
            # build or load summary
            summary_file = os.path.join(
                self.summaries_path, f"{mid}_summary.txt"
            )
            if os.path.exists(summary_file):
                with open(summary_file) as f:
                    return i, f.read()
            # optional: vectorize & retrieve, but here we just summarize full text
            result = chain.invoke({"metadata": text})
            with open(summary_file, "w") as f:
                f.write(result)
            return i, result

        with ThreadPoolExecutor(
            max_workers=min(8, len(state["materials"]))
        ) as exe:
            futures = [
                exe.submit(process, i, m)
                for i, m in enumerate(state["materials"])
            ]
            for future in tqdm(futures, desc="Summarizing materials"):
                i, summ = future.result()
                summaries[i] = summ

        return {**state, "summaries": summaries}

    def _aggregate_node(self, state: dict) -> dict:
        """Combine all summaries into a single, coherent answer."""
        combined = "\n\n----\n\n".join(
            f"[{i + 1}] {m['material_id']}\n\n{summary}"
            for i, (m, summary) in enumerate(
                zip(state["materials"], state["summaries"])
            )
        )

        prompt = ChatPromptTemplate.from_template("""
        You are a materials informatics assistant. Below are brief summaries of several materials:

        {summaries}

        Answer the user’s question in context:

        {context}
                """)
        chain = prompt | self.llm | StrOutputParser()
        final = chain.invoke({
            "summaries": combined,
            "context": state["context"],
        })
        return {**state, "final_summary": final}

    def _build_graph(self):
        self.add_node(self._fetch_node)
        if self.summarize:
            self.add_node(self._summarize_node)
            self.add_node(self._aggregate_node)

            self.graph.set_entry_point("_fetch_node")
            self.graph.add_edge("_fetch_node", "_summarize_node")
            self.graph.add_edge("_summarize_node", "_aggregate_node")
            self.graph.set_finish_point("_aggregate_node")
        else:
            self.graph.set_entry_point("_fetch_node")
            self.graph.set_finish_point("_fetch_node")

optimization_agent

run_cmd(query, state)

Run a commandline command from using the subprocess package in python

Parameters:

Name	Type	Description	Default
query	str	commandline command to be run as a string given to the subprocess.run command.	required

Source code in src/ursa/agents/optimization_agent.py

@tool
def run_cmd(query: str, state: Annotated[dict, InjectedState]) -> str:
    """
    Run a commandline command from using the subprocess package in python

    Args:
        query: commandline command to be run as a string given to the subprocess.run command.
    """
    workspace_dir = state["workspace"]
    LOGGER.info("Running command: %s", query)
    try:
        process = subprocess.Popen(
            query.split(" "),
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            cwd=workspace_dir,
        )

        stdout, stderr = process.communicate(timeout=60000)
        LOGGER.info(
            "Command finished: returncode=%s stdout_chars=%s stderr_chars=%s",
            process.returncode,
            len(stdout or ""),
            len(stderr or ""),
        )
    except KeyboardInterrupt:
        LOGGER.warning("Keyboard interrupt while running command: %s", query)
        stdout, stderr = "", "KeyboardInterrupt:"

    return f"STDOUT: {stdout} and STDERR: {stderr}"

write_code(code, filename, state)

Writes python or Julia code to a file in the given workspace as requested.

Parameters:

Name	Type	Description	Default
code	str	The code to write	required
filename	str	the filename with an appropriate extension for programming language (.py for python, .jl for Julia, etc.)	required

Returns:

Type	Description
str	Execution results

Source code in src/ursa/agents/optimization_agent.py

@tool
def write_code(
    code: str,
    filename: str,
    state: Annotated[dict, InjectedState],
) -> str:
    """
    Writes python or Julia code to a file in the given workspace as requested.

    Args:
        code: The code to write
        filename: the filename with an appropriate extension for programming language (.py for python, .jl for Julia, etc.)

    Returns:
        Execution results
    """
    workspace_dir = state["workspace"]
    try:
        LOGGER.info("Writing file: %s", filename)
        # Extract code if wrapped in markdown code blocks
        if "```" in code:
            code_parts = code.split("```")
            if len(code_parts) >= 3:
                # Extract the actual code
                if "\n" in code_parts[1]:
                    code = "\n".join(code_parts[1].strip().split("\n")[1:])
                else:
                    code = code_parts[2].strip()

        # Write code to a file
        code_file = os.path.join(workspace_dir, filename)

        with open(code_file, "w") as f:
            f.write(code)
        LOGGER.info("File written: %s", code_file)

        return f"File {filename} written successfully."

    except Exception:
        LOGGER.exception("Error generating code for %s", filename)
        # Return minimal code that prints the error
        return f"Failed to write {filename} successfully."

planning_agent

PlanningAgent

Bases: BaseAgent[PlanningState]

Source code in src/ursa/agents/planning_agent.py

class PlanningAgent(BaseAgent[PlanningState]):
    agent_state = PlanningState

    def __init__(
        self,
        llm: BaseChatModel,
        max_reflection_steps: int = 1,
        **kwargs,
    ):
        super().__init__(llm, **kwargs)
        self.planner_prompt = planner_prompt
        self.reflection_prompt = reflection_prompt
        self.max_reflection_steps = max_reflection_steps

    def format_result(self, state: PlanningState) -> str:
        return str(state["plan"])

    def generation_node(
        self,
        state: PlanningState,
        config: RunnableConfig | None = None,
    ) -> PlanningState:
        """
        Plan generation with structured output. Produces a JSON string in messages
        and a parsed list of steps in state["plan_steps"].
        """
        events = self.events(config)
        events.emit("Drafting plan", stage="generate")
        messages = cast(list, state.get("messages"))
        if isinstance(messages[0], SystemMessage):
            messages[0] = SystemMessage(content=self.planner_prompt)
        else:
            messages = [SystemMessage(content=self.planner_prompt)] + messages

        plan = cast(
            Plan,
            invoke_structured(
                self.llm,
                Plan,
                messages,
                context="planning generation",
                repair=3,
            ),
        )
        events.emit(
            "Drafted plan",
            stage="generate_result",
            steps=[step.model_dump() for step in plan.steps],
        )

        return {
            "plan": plan,
            "messages": [AIMessage(content=plan.model_dump_json())],
            "reflection_steps": state.get(
                "reflection_steps", self.max_reflection_steps
            ),
        }

    def reflection_node(
        self,
        state: PlanningState,
        config: RunnableConfig | None = None,
    ) -> PlanningState:
        events = self.events(config)
        events.emit("Reviewing plan", stage="reflect")
        cls_map = {"ai": HumanMessage, "human": AIMessage}
        translated = [state["messages"][0]] + [
            cls_map[msg.type](content=msg.content)
            for msg in state["messages"][1:]
        ]
        translated = [SystemMessage(content=reflection_prompt)] + translated
        res = StrOutputParser().invoke(
            self.llm.invoke(
                translated,
                self.build_config(tags=["planner", "reflect"]),
            )
        )

        if not res.strip():
            # Some providers can return an empty reflection message; treat that as
            # "no objections" so we do not regenerate from an empty human turn.
            res = "[APPROVED]"

        approved = "[APPROVED]" in res
        events.emit(
            "Plan approved" if approved else "Plan needs another pass",
            stage="reflect_result",
            approved=approved,
            reason=res,
        )
        return {
            "plan": state["plan"],
            "messages": [HumanMessage(content=res)],
            "reflection_steps": state["reflection_steps"] - 1,
        }

    def _build_graph(self):
        self.add_node(self.generation_node, "generate")
        self.add_node(self.reflection_node, "reflect")
        self.graph.set_entry_point("generate")
        self.graph.add_conditional_edges(
            "generate",
            self._wrap_cond(
                _should_reflect, "should_reflect", "planning_agent"
            ),
            {"reflect": "reflect", "END": END},
        )
        self.graph.add_conditional_edges(
            "reflect",
            self._wrap_cond(
                _should_regenerate, "should_regenerate", "planning_agent"
            ),
            {"generate": "generate", "END": END},
        )

generation_node(state, config=None)

Plan generation with structured output. Produces a JSON string in messages and a parsed list of steps in state["plan_steps"].

Source code in src/ursa/agents/planning_agent.py

def generation_node(
    self,
    state: PlanningState,
    config: RunnableConfig | None = None,
) -> PlanningState:
    """
    Plan generation with structured output. Produces a JSON string in messages
    and a parsed list of steps in state["plan_steps"].
    """
    events = self.events(config)
    events.emit("Drafting plan", stage="generate")
    messages = cast(list, state.get("messages"))
    if isinstance(messages[0], SystemMessage):
        messages[0] = SystemMessage(content=self.planner_prompt)
    else:
        messages = [SystemMessage(content=self.planner_prompt)] + messages

    plan = cast(
        Plan,
        invoke_structured(
            self.llm,
            Plan,
            messages,
            context="planning generation",
            repair=3,
        ),
    )
    events.emit(
        "Drafted plan",
        stage="generate_result",
        steps=[step.model_dump() for step in plan.steps],
    )

    return {
        "plan": plan,
        "messages": [AIMessage(content=plan.model_dump_json())],
        "reflection_steps": state.get(
            "reflection_steps", self.max_reflection_steps
        ),
    }

PlanningState

Bases: TypedDict

State dictionary for planning agent

Source code in src/ursa/agents/planning_agent.py

class PlanningState(TypedDict, total=False):
    """State dictionary for planning agent"""

    plan: Plan
    messages: Annotated[list, add_messages]
    reflection_steps: int

prompting_agent

PromptingAgent

Bases: BaseAgent[PromptingState]

Prompting Agent

Iterates with the user to turn an initial rough prompt into clear, self-contained instructions for a downstream agentic workflow. The agent proposes an improved prompt, accepts human feedback in subsequent turns, and marks the result approved when the user confirms it.

Source code in src/ursa/agents/prompting_agent.py

class PromptingAgent(BaseAgent[PromptingState]):
    """Prompting Agent

    Iterates with the user to turn an initial rough prompt into clear,
    self-contained instructions for a downstream agentic workflow. The agent
    proposes an improved prompt, accepts human feedback in subsequent turns, and
    marks the result approved when the user confirms it.
    """

    state_type = PromptingState

    def __init__(
        self,
        llm: BaseChatModel,
        use_web: bool = False,
        extra_execution_tools: list[BaseTool] | None = None,
        **kwargs,
    ):
        super().__init__(llm, **kwargs)
        tool_context = _format_available_tool_context(
            use_web=use_web,
            extra_execution_tools=extra_execution_tools,
        )
        self.prompting_agent_prompt = (
            prompting_agent_prompt
            + "\n\nAvailable downstream tool context:\n\n"
            + tool_context
        )

    def format_query(
        self, prompt: str, state: PromptingState | None = None
    ) -> PromptingState:
        """Format a user message into the prompting-agent state.

        On the first turn, the user message is treated as the original prompt to
        refine. On later turns, it is treated as either approval or feedback on
        the previously proposed prompt.
        """
        if state is None:
            return PromptingState(
                messages=[HumanMessage(content=prompt)],
                approved=False,
            )

        state["messages"].append(HumanMessage(content=prompt))
        state["approved"] = False
        return state

    def format_result(self, state: PromptingState) -> str:
        """Return the current refined prompt as plain text."""
        prompt = state.get("prompt")
        if prompt:
            if state.get("approved"):
                return f"Approved prompt:\n\n{prompt}"
            return prompt
        if state.get("messages"):
            return state["messages"][-1].text
        return ""

    def route_node(self, state: PromptingState) -> PromptingState:
        """No-op entry node used to route approval before invoking the LLM."""
        return {}

    def proposal_node(self, state: PromptingState) -> PromptingState:
        """Generate or revise a downstream-agent prompt."""
        print("PromptingAgent: refining prompt . . .")

        new_state, full_overwrite = self.prepare_messages_context(
            state,
            system_prompt=self.prompting_agent_prompt,
            # This agent only sees tool descriptions as prompt context. It does
            # not bind or call tools, so dangling-tool patching is unnecessary.
            patch_dangling=False,
        )

        messages = cast(list, new_state.get("messages", []))
        if messages and isinstance(messages[0], SystemMessage):
            messages[0] = SystemMessage(content=self.prompting_agent_prompt)
        else:
            messages = [
                SystemMessage(content=self.prompting_agent_prompt)
            ] + messages

        response = self.llm.invoke(
            messages,
            self.build_config(tags=["prompting", "refine"]),
        )
        proposed_prompt = response.text

        return self.messages_update(
            new_state,
            [AIMessage(content=proposed_prompt)],
            full_overwrite=full_overwrite,
            extra={"prompt": proposed_prompt, "approved": False},
        )

    def approval_node(self, state: PromptingState) -> PromptingState:
        """Record human approval of the current prompt.

        Human review happens outside the graph: the user either invokes the agent
        again with feedback, causing another proposal cycle, or replies with a
        clear approval phrase, causing the existing prompt to be finalized.
        """
        prompt = state.get("prompt", "")
        print("PromptingAgent: prompt approved")
        return {
            "prompt": prompt,
            "approved": True,
            "messages": [AIMessage(content=f"Approved prompt:\n\n{prompt}")],
        }

    def _build_graph(self):
        self.add_node(self.route_node, "route")
        self.add_node(self.proposal_node, "propose")
        self.add_node(self.approval_node, "approve")
        self.graph.set_entry_point("route")
        self.graph.add_conditional_edges(
            "route",
            self._wrap_cond(
                _should_mark_approved,
                "should_mark_approved",
                "prompting_agent",
            ),
            {"approve": "approve", "propose": "propose"},
        )
        self.graph.add_edge("propose", END)
        self.graph.set_finish_point("approve")

approval_node(state)

Record human approval of the current prompt.

Human review happens outside the graph: the user either invokes the agent again with feedback, causing another proposal cycle, or replies with a clear approval phrase, causing the existing prompt to be finalized.

Source code in src/ursa/agents/prompting_agent.py

def approval_node(self, state: PromptingState) -> PromptingState:
    """Record human approval of the current prompt.

    Human review happens outside the graph: the user either invokes the agent
    again with feedback, causing another proposal cycle, or replies with a
    clear approval phrase, causing the existing prompt to be finalized.
    """
    prompt = state.get("prompt", "")
    print("PromptingAgent: prompt approved")
    return {
        "prompt": prompt,
        "approved": True,
        "messages": [AIMessage(content=f"Approved prompt:\n\n{prompt}")],
    }

format_query(prompt, state=None)

Format a user message into the prompting-agent state.

On the first turn, the user message is treated as the original prompt to refine. On later turns, it is treated as either approval or feedback on the previously proposed prompt.

Source code in src/ursa/agents/prompting_agent.py

def format_query(
    self, prompt: str, state: PromptingState | None = None
) -> PromptingState:
    """Format a user message into the prompting-agent state.

    On the first turn, the user message is treated as the original prompt to
    refine. On later turns, it is treated as either approval or feedback on
    the previously proposed prompt.
    """
    if state is None:
        return PromptingState(
            messages=[HumanMessage(content=prompt)],
            approved=False,
        )

    state["messages"].append(HumanMessage(content=prompt))
    state["approved"] = False
    return state

format_result(state)

Return the current refined prompt as plain text.

Source code in src/ursa/agents/prompting_agent.py

def format_result(self, state: PromptingState) -> str:
    """Return the current refined prompt as plain text."""
    prompt = state.get("prompt")
    if prompt:
        if state.get("approved"):
            return f"Approved prompt:\n\n{prompt}"
        return prompt
    if state.get("messages"):
        return state["messages"][-1].text
    return ""

proposal_node(state)

Generate or revise a downstream-agent prompt.

Source code in src/ursa/agents/prompting_agent.py

def proposal_node(self, state: PromptingState) -> PromptingState:
    """Generate or revise a downstream-agent prompt."""
    print("PromptingAgent: refining prompt . . .")

    new_state, full_overwrite = self.prepare_messages_context(
        state,
        system_prompt=self.prompting_agent_prompt,
        # This agent only sees tool descriptions as prompt context. It does
        # not bind or call tools, so dangling-tool patching is unnecessary.
        patch_dangling=False,
    )

    messages = cast(list, new_state.get("messages", []))
    if messages and isinstance(messages[0], SystemMessage):
        messages[0] = SystemMessage(content=self.prompting_agent_prompt)
    else:
        messages = [
            SystemMessage(content=self.prompting_agent_prompt)
        ] + messages

    response = self.llm.invoke(
        messages,
        self.build_config(tags=["prompting", "refine"]),
    )
    proposed_prompt = response.text

    return self.messages_update(
        new_state,
        [AIMessage(content=proposed_prompt)],
        full_overwrite=full_overwrite,
        extra={"prompt": proposed_prompt, "approved": False},
    )

route_node(state)

No-op entry node used to route approval before invoking the LLM.

Source code in src/ursa/agents/prompting_agent.py

def route_node(self, state: PromptingState) -> PromptingState:
    """No-op entry node used to route approval before invoking the LLM."""
    return {}

PromptingState

Bases: TypedDict

State dictionary for the prompting agent.

The agent keeps the full refinement conversation in messages and stores the latest proposed downstream-agent prompt in prompt. approved indicates whether the human has accepted the current prompt.

Source code in src/ursa/agents/prompting_agent.py

class PromptingState(TypedDict, total=False):
    """State dictionary for the prompting agent.

    The agent keeps the full refinement conversation in ``messages`` and stores the
    latest proposed downstream-agent prompt in ``prompt``. ``approved`` indicates
    whether the human has accepted the current prompt.
    """

    messages: Annotated[list, add_messages]
    prompt: str
    approved: bool