ursa-ai 0.4.2__py3-none-any.whl → 0.6.0__py3-none-any.whl

ursa/agents/execution_agent.py

@@ -3,7 +3,7 @@ import os
 # from langchain_core.runnables.graph import MermaidDrawMethod
 import subprocess
 from pathlib import Path
-from typing import Annotated, Any, Literal, Optional
+from typing import Annotated, Any, Literal, Mapping, Optional
 
 import randomname
 from langchain_community.tools import (
@@ -17,11 +17,11 @@ from langchain_core.messages import (
     ToolMessage,
 )
 from langchain_core.tools import InjectedToolCallId, tool
-from langgraph.graph import END, START, StateGraph
+from langgraph.graph import StateGraph
 from langgraph.graph.message import add_messages
 from langgraph.prebuilt import InjectedState, ToolNode
 from langgraph.types import Command
-from litellm import ContentPolicyViolationError
+from litellm.exceptions import ContentPolicyViolationError
 
 # Rich
 from rich import get_console
@@ -29,7 +29,11 @@ from rich.panel import Panel
 from rich.syntax import Syntax
 from typing_extensions import TypedDict
 
-from ..prompt_library.execution_prompts import executor_prompt, summarize_prompt
+from ..prompt_library.execution_prompts import (
+    executor_prompt,
+    safety_prompt,
+    summarize_prompt,
+)
 from ..util.diff_renderer import DiffRenderer
 from ..util.memory_logger import AgentMemory
 from .base import BaseAgent
@@ -62,6 +66,7 @@ class ExecutionAgent(BaseAgent):
     ):
         super().__init__(llm, **kwargs)
         self.agent_memory = agent_memory
+        self.safety_prompt = safety_prompt
         self.executor_prompt = executor_prompt
         self.summarize_prompt = summarize_prompt
         self.tools = [run_cmd, write_code, edit_code, search_tool]
@@ -69,7 +74,7 @@ class ExecutionAgent(BaseAgent):
         self.llm = self.llm.bind_tools(self.tools)
         self.log_state = log_state
 
-        self._initialize_agent()
+        self._action = self._build_graph()
 
     # Define the function that calls the model
     def query_executor(self, state: ExecutionState) -> ExecutionState:
@@ -115,8 +120,7 @@ class ExecutionAgent(BaseAgent):
             ] + state["messages"]
         try:
             response = self.llm.invoke(
-                new_state["messages"],
-                {"configurable": {"thread_id": self.thread_id}},
+                new_state["messages"], self.build_config(tags=["agent"])
             )
         except ContentPolicyViolationError as e:
             print("Error: ", e, " ", new_state["messages"][-1].content)
@@ -129,7 +133,7 @@ class ExecutionAgent(BaseAgent):
         messages = [SystemMessage(content=summarize_prompt)] + state["messages"]
         try:
             response = self.llm.invoke(
-                messages, {"configurable": {"thread_id": self.thread_id}}
+                messages, self.build_config(tags=["summarize"])
             )
         except ContentPolicyViolationError as e:
             print("Error: ", e, " ", messages[-1].content)
@@ -181,11 +185,8 @@ class ExecutionAgent(BaseAgent):
             if call_name == "run_cmd":
                 query = tool_call["args"]["query"]
                 safety_check = self.llm.invoke(
-                    (
-                        "Assume commands to run/install python and Julia files are safe because "
-                        "the files are from a trusted source. "
-                        f"Explain why, followed by an answer [YES] or [NO]. Is this command safe to run: {query}"
-                    )
+                    self.safety_prompt + query,
+                    self.build_config(tags=["safety_check"]),
                 )
 
                 if "[NO]" in safety_check.content:
@@ -222,53 +223,90 @@ class ExecutionAgent(BaseAgent):
 
         return new_state
 
-    def _initialize_agent(self):
-        self.graph = StateGraph(ExecutionState)
+    def _build_graph(self):
+        graph = StateGraph(ExecutionState)
 
-        self.graph.add_node("agent", self.query_executor)
-        self.graph.add_node("action", self.tool_node)
-        self.graph.add_node("summarize", self.summarize)
-        self.graph.add_node("safety_check", self.safety_check)
+        self.add_node(graph, self.query_executor, "agent")
+        self.add_node(graph, self.tool_node, "action")
+        self.add_node(graph, self.summarize, "summarize")
+        self.add_node(graph, self.safety_check, "safety_check")
 
         # Set the entrypoint as `agent`
         # This means that this node is the first one called
-        self.graph.add_edge(START, "agent")
+        graph.set_entry_point("agent")
 
-        self.graph.add_conditional_edges(
+        graph.add_conditional_edges(
             "agent",
-            should_continue,
-            {
-                "continue": "safety_check",
-                "summarize": "summarize",
-            },
+            self._wrap_cond(should_continue, "should_continue", "execution"),
+            {"continue": "safety_check", "summarize": "summarize"},
         )
 
-        self.graph.add_conditional_edges(
+        graph.add_conditional_edges(
             "safety_check",
-            command_safe,
-            {
-                "safe": "action",
-                "unsafe": "agent",
-            },
+            self._wrap_cond(command_safe, "command_safe", "execution"),
+            {"safe": "action", "unsafe": "agent"},
         )
 
-        self.graph.add_edge("action", "agent")
-        self.graph.add_edge("summarize", END)
+        graph.add_edge("action", "agent")
+        graph.set_finish_point("summarize")
 
-        self.action = self.graph.compile(checkpointer=self.checkpointer)
+        return graph.compile(checkpointer=self.checkpointer)
         # self.action.get_graph().draw_mermaid_png(output_file_path="execution_agent_graph.png", draw_method=MermaidDrawMethod.PYPPETEER)
 
-    def run(self, prompt, recursion_limit=1000):
-        inputs = {"messages": [HumanMessage(content=prompt)]}
-        return self.action.invoke(
-            inputs,
-            {
-                "recursion_limit": recursion_limit,
-                "configurable": {"thread_id": self.thread_id},
-            },
+    def _invoke(
+        self, inputs: Mapping[str, Any], recursion_limit: int = 999_999, **_
+    ):
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
+        return self._action.invoke(inputs, config)
+
+    # this is trying to stop people bypassing invoke
+    @property
+    def action(self):
+        raise AttributeError(
+            "Use .stream(...) or .invoke(...); direct .action access is unsupported."
         )
 
 
+def _snip_text(text: str, max_chars: int) -> tuple[str, bool]:
+    if text is None:
+        return "", False
+    if max_chars <= 0:
+        return "", len(text) > 0
+    if len(text) <= max_chars:
+        return text, False
+    head = max_chars // 2
+    tail = max_chars - head
+    return (
+        text[:head]
+        + f"\n... [snipped {len(text) - max_chars} chars] ...\n"
+        + text[-tail:],
+        True,
+    )
+
+
+def _fit_streams_to_budget(stdout: str, stderr: str, total_budget: int):
+    label_overhead = len("STDOUT:\n") + len("\nSTDERR:\n")
+    budget = max(0, total_budget - label_overhead)
+
+    if len(stdout) + len(stderr) <= budget:
+        return stdout, stderr
+
+    total_len = max(1, len(stdout) + len(stderr))
+    stdout_budget = int(budget * (len(stdout) / total_len))
+    stderr_budget = budget - stdout_budget
+
+    stdout_snip, _ = _snip_text(stdout, stdout_budget)
+    stderr_snip, _ = _snip_text(stderr, stderr_budget)
+    return stdout_snip, stderr_snip
+
+
+# the idea here is that we just set a limit - the user could overload
+# that in their env, or maybe we could pull this out of the LLM parameters
+MAX_TOOL_MSG_CHARS = int(os.getenv("MAX_TOOL_MSG_CHARS", "50000"))
+
+
 @tool
 def run_cmd(query: str, state: Annotated[dict, InjectedState]) -> str:
     """
@@ -293,10 +331,15 @@ def run_cmd(query: str, state: Annotated[dict, InjectedState]) -> str:
         print("Keyboard Interrupt of command: ", query)
         stdout, stderr = "", "KeyboardInterrupt:"
 
-    print("STDOUT: ", stdout)
-    print("STDERR: ", stderr)
+    # Fit BOTH streams under a single overall cap
+    stdout_fit, stderr_fit = _fit_streams_to_budget(
+        stdout or "", stderr or "", MAX_TOOL_MSG_CHARS
+    )
+
+    print("STDOUT: ", stdout_fit)
+    print("STDERR: ", stderr_fit)
 
-    return f"STDOUT: {stdout} and STDERR: {stderr}"
+    return f"STDOUT:\n{stdout_fit}\nSTDERR:\n{stderr_fit}"
 
 
 def _strip_fences(snippet: str) -> str:
@@ -499,8 +542,9 @@ def main():
     inputs = {
         "messages": [HumanMessage(content=problem_string)]
     }  # , "workspace":"dummy_test"}
-    result = execution_agent.action.invoke(
-        inputs, {"configurable": {"thread_id": execution_agent.thread_id}}
+    result = execution_agent.invoke(
+        inputs,
+        config={"configurable": {"thread_id": execution_agent.thread_id}},
     )
     print(result["messages"][-1].content)
     return result

ursa/agents/hypothesizer_agent.py

@@ -3,12 +3,12 @@ import ast
 # from langchain_community.tools import TavilySearchResults
 # from textwrap                  import dedent
 from datetime import datetime
-from typing import List, Literal, TypedDict
+from typing import Any, List, Literal, Mapping, TypedDict
 
 from langchain_community.tools import DuckDuckGoSearchResults
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage
-from langgraph.graph import END, StateGraph
+from langgraph.graph import StateGraph
 
 from ..prompt_library.hypothesizer_prompts import (
     competitor_prompt,
@@ -53,7 +53,7 @@ class HypothesizerAgent(BaseAgent):
         #     max_results=10, search_depth="advanced", include_answer=False
         # )
 
-        self._initialize_agent()
+        self._action = self._build_graph()
 
     def agent1_generate_solution(
         self, state: HypothesizerState
@@ -444,68 +444,65 @@ class HypothesizerAgent(BaseAgent):
         )
         return new_state
 
-    def _initialize_agent(self):
+    def _build_graph(self):
         # Initialize the graph
-        self.graph = StateGraph(HypothesizerState)
+        graph = StateGraph(HypothesizerState)
 
         # Add nodes
-        self.graph.add_node("agent1", self.agent1_generate_solution)
-        self.graph.add_node("agent2", self.agent2_critique)
-        self.graph.add_node("agent3", self.agent3_competitor_perspective)
-        self.graph.add_node("increment_iteration", self.increment_iteration)
-        self.graph.add_node("finalize", self.generate_solution)
-        self.graph.add_node("print_sites", self.print_visited_sites)
-        self.graph.add_node(
-            "summarize_as_latex", self.summarize_process_as_latex
+        self.add_node(graph, self.agent1_generate_solution, "agent1")
+        self.add_node(graph, self.agent2_critique, "agent2")
+        self.add_node(graph, self.agent3_competitor_perspective, "agent3")
+        self.add_node(graph, self.increment_iteration, "increment_iteration")
+        self.add_node(graph, self.generate_solution, "finalize")
+        self.add_node(graph, self.print_visited_sites, "print_sites")
+        self.add_node(
+            graph, self.summarize_process_as_latex, "summarize_as_latex"
         )
         # self.graph.add_node("compile_pdf",                compile_summary_to_pdf)
 
         # Add simple edges for the known flow
-        self.graph.add_edge("agent1", "agent2")
-        self.graph.add_edge("agent2", "agent3")
-        self.graph.add_edge("agent3", "increment_iteration")
+        graph.add_edge("agent1", "agent2")
+        graph.add_edge("agent2", "agent3")
+        graph.add_edge("agent3", "increment_iteration")
 
         # Then from increment_iteration, we have a conditional:
         # If we 'continue', we go back to agent1
         # If we 'finish', we jump to the finalize node
-        self.graph.add_conditional_edges(
+        graph.add_conditional_edges(
             "increment_iteration",
             should_continue,
             {"continue": "agent1", "finish": "finalize"},
         )
 
-        self.graph.add_edge("finalize", "summarize_as_latex")
-        self.graph.add_edge("summarize_as_latex", "print_sites")
-        self.graph.add_edge("print_sites", END)
+        graph.add_edge("finalize", "summarize_as_latex")
+        graph.add_edge("summarize_as_latex", "print_sites")
         # self.graph.add_edge("summarize_as_latex", "compile_pdf")
         # self.graph.add_edge("compile_pdf", "print_sites")
 
         # Set the entry point
-        self.graph.set_entry_point("agent1")
+        graph.set_entry_point("agent1")
+        graph.set_finish_point("print_sites")
 
-        self.action = self.graph.compile(checkpointer=self.checkpointer)
+        return graph.compile(checkpointer=self.checkpointer)
         # self.action.get_graph().draw_mermaid_png(output_file_path="hypothesizer_agent_graph.png", draw_method=MermaidDrawMethod.PYPPETEER)
 
-    def run(self, prompt, max_iter=3, recursion_limit=99999):
-        # Initialize the state
-        initial_state = HypothesizerState(
-            question=prompt,
-            current_iteration=0,
-            max_iterations=max_iter,
-            agent1_solution=[],
-            agent2_critiques=[],
-            agent3_perspectives=[],
-            solution="",
+    def _invoke(
+        self, inputs: Mapping[str, Any], recursion_limit: int = 100000, **_
+    ):
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
         )
-        # Run the graph
-        result = self.action.invoke(
-            initial_state,
-            {
-                "recursion_limit": recursion_limit,
-                "configurable": {"thread_id": self.thread_id},
-            },
-        )
-        return result["solution"]
+        if "prompt" not in inputs:
+            raise KeyError("'prompt' is a required arguments")
+
+        inputs["max_iterations"] = inputs.get("max_iterations", 3)
+        inputs["current_iteration"] = 0
+        inputs["agent1_solution"] = []
+        inputs["agent2_critiques"] = []
+        inputs["agent3_perspectives"] = []
+        inputs["solution"] = ""
+
+        return self._action.invoke(inputs, config)
 
 
 def should_continue(state: HypothesizerState) -> Literal["continue", "finish"]:
@@ -583,7 +580,7 @@ if __name__ == "__main__":
 
     print("[DEBUG] Invoking the graph...")
     # Run the graph
-    result = hypothesizer_agent.action.invoke(
+    result = hypothesizer_agent.invoke(
         initial_state,
         {
             "recursion_limit": 999999,

ursa/agents/lammps_agent.py

@@ -1,17 +1,22 @@
 import json
 import os
 import subprocess
-from typing import Any, Dict, List, Optional, TypedDict
+from typing import Any, Dict, List, Mapping, Optional, TypedDict
 
-import atomman as am
 import tiktoken
-import trafilatura
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 from langgraph.graph import END, StateGraph
 
 from .base import BaseAgent
 
+working = True
+try:
+    import atomman as am
+    import trafilatura
+except Exception:
+    working = False
+
 
 class LammpsState(TypedDict, total=False):
     simulation_task: str
@@ -50,6 +55,10 @@ class LammpsAgent(BaseAgent):
         max_tokens: int = 200000,
         **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]' ."
+            )
         self.max_potentials = max_potentials
         self.max_fix_attempts = max_fix_attempts
         self.mpi_procs = mpi_procs
@@ -144,7 +153,7 @@ class LammpsAgent(BaseAgent):
             | self.str_parser
         )
 
-        self.graph = self._build_graph().compile()
+        self._action = self._build_graph()
 
     @staticmethod
     def _safe_json_loads(s: str) -> Dict[str, Any]:
@@ -340,53 +349,66 @@ class LammpsAgent(BaseAgent):
     def _build_graph(self):
         g = StateGraph(LammpsState)
 
-        g.add_node("find_potentials", self._find_potentials)
-        g.add_node("summarize_one", self._summarize_one)
-        g.add_node("build_summaries", self._build_summaries)
-        g.add_node("choose", self._choose)
-        g.add_node("author", self._author)
-        g.add_node("run_lammps", self._run_lammps)
-        g.add_node("fix", self._fix)
+        self.add_node(g, self._find_potentials)
+        self.add_node(g, self._summarize_one)
+        self.add_node(g, self._build_summaries)
+        self.add_node(g, self._choose)
+        self.add_node(g, self._author)
+        self.add_node(g, self._run_lammps)
+        self.add_node(g, self._fix)
 
-        g.set_entry_point("find_potentials")
+        g.set_entry_point("_find_potentials")
 
         g.add_conditional_edges(
-            "find_potentials",
+            "_find_potentials",
             self._should_summarize,
             {
-                "summarize_one": "summarize_one",
-                "summarize_done": "build_summaries",
+                "summarize_one": "_summarize_one",
+                "summarize_done": "_build_summaries",
                 "done_no_matches": END,
             },
         )
 
         g.add_conditional_edges(
-            "summarize_one",
+            "_summarize_one",
             self._should_summarize,
             {
-                "summarize_one": "summarize_one",
-                "summarize_done": "build_summaries",
+                "summarize_one": "_summarize_one",
+                "summarize_done": "_build_summaries",
             },
         )
 
-        g.add_edge("build_summaries", "choose")
-        g.add_edge("choose", "author")
-        g.add_edge("author", "run_lammps")
+        g.add_edge("_build_summaries", "_choose")
+        g.add_edge("_choose", "_author")
+        g.add_edge("_author", "_run_lammps")
 
         g.add_conditional_edges(
-            "run_lammps",
+            "_run_lammps",
             self._route_run,
             {
-                "need_fix": "fix",
+                "need_fix": "_fix",
                 "done_success": END,
                 "done_failed": END,
             },
         )
-        g.add_edge("fix", "run_lammps")
-        return g
+        g.add_edge("_fix", "_run_lammps")
+        return g.compile(checkpointer=self.checkpointer)
 
-    def run(self, simulation_task, elements):
-        return self.graph.invoke(
-            {"simulation_task": simulation_task, "elements": elements},
-            {"recursion_limit": 999_999},
+    def _invoke(
+        self,
+        inputs: Mapping[str, Any],
+        *,
+        summarize: bool | None = None,
+        recursion_limit: int = 1000,
+        **_,
+    ) -> str:
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
         )
+
+        if "simulation_task" not in inputs or "elements" not in inputs:
+            raise KeyError(
+                "'simulation_task' and 'elements' are required arguments"
+            )
+
+        return self._action.invoke(inputs, config)

ursa/agents/mp_agent.py

@@ -2,7 +2,7 @@ import json
 import os
 import re
 from concurrent.futures import ThreadPoolExecutor
-from typing import Dict
+from typing import Any, Dict, Mapping
 
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
@@ -50,7 +50,7 @@ class MaterialsProjectAgent(BaseAgent):
         os.makedirs(self.database_path, exist_ok=True)
         os.makedirs(self.summaries_path, exist_ok=True)
 
-        self.graph = self._build_graph()
+        self._action = self._build_graph()
 
     def _fetch_node(self, state: Dict) -> Dict:
         f = state["query"]
@@ -148,31 +148,56 @@ You are a materials-science assistant. Given the following metadata about a mate
         return {**state, "final_summary": final}
 
     def _build_graph(self):
-        g = StateGraph(dict)  # using plain dict for state
-        g.add_node("fetch", self._fetch_node)
+        graph = StateGraph(dict)  # using plain dict for state
+        self.add_node(graph, self._fetch_node)
         if self.summarize:
-            g.add_node("summarize", self._summarize_node)
-            g.add_node("aggregate", self._aggregate_node)
-            g.set_entry_point("fetch")
-            g.add_edge("fetch", "summarize")
-            g.add_edge("summarize", "aggregate")
-            g.set_finish_point("aggregate")
+            self.add_node(graph, self._summarize_node)
+            self.add_node(graph, self._aggregate_node)
+
+            graph.set_entry_point("_fetch_node")
+            graph.add_edge("_fetch_node", "_summarize_node")
+            graph.add_edge("_summarize_node", "_aggregate_node")
+            graph.set_finish_point("_aggregate_node")
         else:
-            g.set_entry_point("fetch")
-            g.set_finish_point("fetch")
-        return g.compile()
+            graph.set_entry_point("_fetch_node")
+            graph.set_finish_point("_fetch_node")
+        return graph.compile(checkpointer=self.checkpointer)
 
-    def run(self, mp_query: str, context: str) -> str:
-        state = {"query": mp_query, "context": context}
-        out = self.graph.invoke(state)
-        if self.summarize:
-            return out.get("final_summary", "")
-        return json.dumps(out.get("materials", []), indent=2)
+    def _invoke(
+        self,
+        inputs: Mapping[str, Any],
+        *,
+        summarize: bool | None = None,
+        recursion_limit: int = 1000,
+        **_,
+    ) -> str:
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
+
+        if "query" not in inputs:
+            if "mp_query" in inputs:
+                # make a shallow copy and rename the key
+                inputs = dict(inputs)
+                inputs["query"] = inputs.pop("mp_query")
+            else:
+                raise KeyError(
+                    "Missing 'query' in inputs (alias 'mp_query' also accepted)."
+                )
+
+        result = self._action.invoke(inputs, config)
+
+        use_summary = self.summarize if summarize is None else summarize
+        return (
+            result.get("final_summary", "No summary generated.")
+            if use_summary
+            else "\n\nFinished Fetching Materials Database Information!"
+        )
 
 
 if __name__ == "__main__":
     agent = MaterialsProjectAgent()
-    resp = agent.run(
+    resp = agent.invoke(
         mp_query="LiFePO4",
         context="What is its band gap and stability, and any synthesis challenges?",
     )