ursa-ai 0.5.0__py3-none-any.whl → 0.6.0rc1__py3-none-any.whl

ursa/agents/optimization_agent.py

@@ -0,0 +1,403 @@
+import os
+import pprint
+import subprocess
+from typing import Annotated, Any, Dict, List, Literal, Mapping
+
+from langchain_community.tools import DuckDuckGoSearchResults
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_core.tools import tool
+from langchain_openai import ChatOpenAI
+from langgraph.graph import END, START, StateGraph
+from langgraph.prebuilt import InjectedState
+from typing_extensions import TypedDict
+
+from ..prompt_library.optimization_prompts import (
+    code_generator_prompt,
+    discretizer_prompt,
+    explainer_prompt,
+    extractor_prompt,
+    feasibility_prompt,
+    math_formulator_prompt,
+    solver_selector_prompt,
+    verifier_prompt,
+)
+from ..tools.feasibility_tools import feasibility_check_auto as fca
+from ..util.helperFunctions import extract_tool_calls, run_tool_calls
+from ..util.optimization_schema import ProblemSpec, SolverSpec
+from .base import BaseAgent
+
+# --- ANSI color codes ---
+GREEN = "\033[92m"
+BLUE = "\033[94m"
+RED = "\033[91m"
+RESET = "\033[0m"
+BOLD = "\033[1m"
+
+
+class OptimizerState(TypedDict):
+    user_input: str
+    problem: str
+    problem_spec: ProblemSpec
+    solver: SolverSpec
+    code: str
+    problem_diagnostic: List[Dict]
+    summary: str
+
+
+class OptimizationAgent(BaseAgent):
+    def __init__(self, llm="OpenAI/gpt-4o", *args, **kwargs):
+        super().__init__(llm, *args, **kwargs)
+        self.extractor_prompt = extractor_prompt
+        self.explainer_prompt = explainer_prompt
+        self.verifier_prompt = verifier_prompt
+        self.code_generator_prompt = code_generator_prompt
+        self.solver_selector_prompt = solver_selector_prompt
+        self.math_formulator_prompt = math_formulator_prompt
+        self.discretizer_prompt = discretizer_prompt
+        self.feasibility_prompt = feasibility_prompt
+        self.tools = [fca]  # [run_cmd, write_code, search_tool, fca]
+        self.llm = self.llm.bind_tools(self.tools)
+        self.tool_maps = {
+            (getattr(t, "name", None) or getattr(t, "__name__", None)): t
+            for i, t in enumerate(self.tools)
+        }
+
+        self._action = self._build_graph()
+
+    # Define the function that calls the model
+    def extractor(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+        new_state["problem"] = self.llm.invoke([
+            SystemMessage(content=self.extractor_prompt),
+            HumanMessage(content=new_state["user_input"]),
+        ]).content
+
+        new_state["problem_diagnostic"] = []
+
+        print("Extractor:\n")
+        pprint.pprint(new_state["problem"])
+        return new_state
+
+    def formulator(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+
+        llm_out = self.llm.with_structured_output(
+            ProblemSpec, include_raw=True
+        ).invoke([
+            SystemMessage(content=self.math_formulator_prompt),
+            HumanMessage(content=state["problem"]),
+        ])
+        new_state["problem_spec"] = llm_out["parsed"]
+        new_state["problem_diagnostic"].extend(
+            extract_tool_calls(llm_out["raw"])
+        )
+
+        print("Formulator:\n")
+        pprint.pprint(new_state["problem_spec"])
+        return new_state
+
+    def discretizer(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+
+        llm_out = self.llm.with_structured_output(
+            ProblemSpec, include_raw=True
+        ).invoke([
+            SystemMessage(content=self.discretizer_prompt),
+            HumanMessage(content=state["problem"]),
+        ])
+        new_state["problem_spec"] = llm_out["parsed"]
+        new_state["problem_diagnostic"].extend(
+            extract_tool_calls(llm_out["raw"])
+        )
+
+        print("Discretizer:\n")
+        pprint.pprint(new_state["problem_spec"])
+
+        return new_state
+
+    def tester(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+
+        llm_out = self.llm.bind(tool_choice="required").invoke([
+            SystemMessage(content=self.feasibility_prompt),
+            HumanMessage(content=str(state["code"])),
+        ])
+
+        tool_log = run_tool_calls(llm_out, self.tool_maps)
+        new_state["problem_diagnostic"].extend(tool_log)
+
+        print("Feasibility Tester:\n")
+        for msg in new_state["problem_diagnostic"]:
+            msg.pretty_print()
+        return new_state
+
+    def selector(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+
+        llm_out = self.llm.with_structured_output(
+            SolverSpec, include_raw=True
+        ).invoke([
+            SystemMessage(content=self.solver_selector_prompt),
+            HumanMessage(content=str(state["problem_spec"])),
+        ])
+        new_state["solver"] = llm_out["parsed"]
+
+        print("Selector:\n ")
+        pprint.pprint(new_state["solver"])
+        return new_state
+
+    def generator(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+
+        new_state["code"] = self.llm.invoke([
+            SystemMessage(content=self.code_generator_prompt),
+            HumanMessage(content=str(state["problem_spec"])),
+        ]).content
+
+        print("Generator:\n")
+        pprint.pprint(new_state["code"])
+        return new_state
+
+    def verifier(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+
+        llm_out = self.llm.with_structured_output(
+            ProblemSpec, include_raw=True
+        ).invoke([
+            SystemMessage(content=self.verifier_prompt),
+            HumanMessage(content=str(state["problem_spec"]) + state["code"]),
+        ])
+        new_state["problem_spec"] = llm_out["parsed"]
+        if hasattr(llm_out, "tool_calls"):
+            tool_log = run_tool_calls(llm_out, self.tool_maps)
+            new_state["problem_diagnostic"].extend(tool_log)
+
+        print("Verifier:\n ")
+        pprint.pprint(new_state["problem_spec"])
+        return new_state
+
+    def explainer(self, state: OptimizerState) -> OptimizerState:
+        new_state = state.copy()
+
+        new_state["summary"] = self.llm.invoke([
+            SystemMessage(content=self.explainer_prompt),
+            HumanMessage(content=state["problem"] + str(state["problem_spec"])),
+            *state["problem_diagnostic"],
+        ]).content
+
+        print("Summary:\n")
+        pprint.pprint(new_state["summary"])
+        return new_state
+
+    def _build_graph(self):
+        graph = StateGraph(OptimizerState)
+
+        self.add_node(graph, self.extractor, "Problem Extractor")
+        self.add_node(graph, self.formulator, "Math Formulator")
+        self.add_node(graph, self.selector, "Solver Selector")
+        self.add_node(graph, self.generator, "Code Generator")
+        self.add_node(graph, self.verifier, "Verifier")
+        self.add_node(graph, self.explainer, "Explainer")
+        self.add_node(graph, self.tester, "Feasibility Tester")
+        self.add_node(graph, self.discretizer, "Discretizer")
+
+        graph.add_edge(START, "Problem Extractor")
+        graph.add_edge("Problem Extractor", "Math Formulator")
+        graph.add_conditional_edges(
+            "Math Formulator",
+            should_discretize,
+            {"discretize": "Discretizer", "continue": "Solver Selector"},
+        )
+        graph.add_edge("Discretizer", "Solver Selector")
+        graph.add_edge("Solver Selector", "Code Generator")
+        graph.add_edge("Code Generator", "Feasibility Tester")
+        graph.add_edge("Feasibility Tester", "Verifier")
+        graph.add_conditional_edges(
+            "Verifier",
+            should_continue,
+            {"continue": "Explainer", "error": "Problem Extractor"},
+        )
+        graph.add_edge("Explainer", END)
+
+        return graph.compile()
+
+    def _invoke(
+        self, inputs: Mapping[str, Any], recursion_limit: int = 100000, **_
+    ):
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
+        if "user_input" not in inputs:
+            try:
+                inputs["user_input"] = inputs["messages"][0].content
+            except KeyError:
+                raise ("'user_input' is a required argument")
+
+        return self._action.invoke(inputs, config)
+
+
+#########  try:
+#########      png_bytes = compiled_graph.get_graph().draw_mermaid_png()
+#########      img = mpimg.imread(io.BytesIO(png_bytes), format='png')  # decode bytes -> array
+
+#########      plt.imshow(img)
+#########      plt.axis('off')
+#########      plt.show()
+#########  except Exception as e:
+#########      # This requires some extra dependencies and is optional
+#########      print(e)
+#########      pass
+
+
+@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"]
+    print("RUNNING: ", query)
+    try:
+        process = subprocess.Popen(
+            query.split(" "),
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            cwd=workspace_dir,
+        )
+
+        stdout, stderr = process.communicate(timeout=60000)
+    except KeyboardInterrupt:
+        print("Keyboard Interrupt of command: ", query)
+        stdout, stderr = "", "KeyboardInterrupt:"
+
+    print("STDOUT: ", stdout)
+    print("STDERR: ", stderr)
+
+    return f"STDOUT: {stdout} and STDERR: {stderr}"
+
+
+@tool
+def write_code(code: str, filename: str, state: Annotated[dict, InjectedState]):
+    """
+    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"]
+    print("Writing filename ", filename)
+    try:
+        # 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)
+        print(f"Written code to file: {code_file}")
+
+        return f"File {filename} written successfully."
+
+    except Exception as e:
+        print(f"Error generating code: {str(e)}")
+        # Return minimal code that prints the error
+        return f"Failed to write {filename} successfully."
+
+
+search_tool = DuckDuckGoSearchResults(output_format="json", num_results=10)
+# search_tool = TavilySearchResults(max_results=10, search_depth="advanced", include_answer=True)
+
+
+# A function to test if discretization is needed
+def should_discretize(
+    state: OptimizerState,
+) -> Literal["Discretize", "continue"]:
+    cons = state["problem_spec"]["constraints"]
+    decs = state["problem_spec"]["decision_variables"]
+
+    if any("infinite-dimensional" in t["tags"] for t in cons) or any(
+        "infinite-dimensional" in t["type"] for t in decs
+    ):
+        # print(f"Problem has infinite-dimensional constraints/decision variables. Needs to be discretized")
+        return "discretize"
+
+    return "continue"
+
+
+# Define the function that determines whether to continue or not
+def should_continue(state: OptimizerState) -> Literal["error", "continue"]:
+    spec = state["problem_spec"]
+    try:
+        status = spec["status"].lower()
+    except KeyError:
+        status = spec["spec"]["status"].lower()
+    if "VERIFIED".lower() in status:
+        return "continue"
+    # Otherwise if there is, we continue
+    else:
+        return "error"
+
+
+def main():
+    model = ChatOpenAI(
+        model="gpt-4o", max_tokens=10000, timeout=None, max_retries=2
+    )
+    execution_agent = OptimizationAgent(llm=model)
+    # execution_agent = execution_agent.bind_tools(feasibility_checker)
+    problem_string = """
+    Solve the following optimal power flow problem
+    System topology and data:
+        - Three buses (nodes) labeled 1, 2 and 3.
+        - One generator at each bus; each can only inject power (no negative output).
+        - Loads of 1 p.u. at bus 1, 2 p.u. at bus 2, and 4 p.u. at bus 3.
+        - Transmission lines connecting every pair of buses, with susceptances (B):
+            - Line 1–2: B₁₂ = 10
+            - Line 1–3: B₁₃ = 20
+            - Line 2–3: B₂₃ = 30
+
+    Decision variables:
+        - Voltage angles θ₁, θ₂, θ₃ (in radians) at buses 1–3.
+        - Generator outputs Pᵍ₁, Pᵍ₂, Pᵍ₃ ≥ 0 (in per-unit).
+
+    Reference angle:
+        - To fix the overall angle‐shift ambiguity, we set θ₁ = 0 (“slack” or reference bus).
+
+    Objective:
+        - Minimize total generation cost with
+            - 𝑐1 = 1
+            - 𝑐2 = 10
+            - 𝑐3 = 100
+
+    Line‐flow limits
+        - Lines 1-2 and 1-3 are thermal‐limited to ±0.5 p.u., line 2-3 is unconstrained.
+    
+    In words:
+    We choose how much each generator should produce (at non-negative cost) and the voltage angles at each bus (with bus 1 set to zero) so that supply exactly meets demand, flows on the critical lines don’t exceed their limits, and the total cost is as small as possible. 
+    Use the tools at your disposal to check if your formulation is feasible.
+    """
+    inputs = {"user_input": problem_string}
+    result = execution_agent.invoke(inputs)
+    print(result["messages"][-1].content)
+    return result
+
+
+if __name__ == "__main__":
+    main()
+
+
+#         min⁡ 𝑃𝑔  𝑐1*𝑃1 + 𝑐2 * 𝑃2 + 𝑐3 * 𝑃3

ursa/agents/planning_agent.py

@@ -1,10 +1,10 @@
 # from langgraph.checkpoint.memory  import MemorySaver
 # from langchain_core.runnables.graph import MermaidDrawMethod
-from typing import Annotated, Any, Dict, List, Optional
+from typing import Annotated, Any, Dict, Iterator, List, Mapping, Optional
 
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
-from langgraph.graph import END, START, StateGraph
+from langgraph.graph import StateGraph
 from langgraph.graph.message import add_messages
 from pydantic import Field
 from typing_extensions import TypedDict
@@ -36,7 +36,7 @@ class PlanningAgent(BaseAgent):
         self.planner_prompt = planner_prompt
         self.formalize_prompt = formalize_prompt
         self.reflection_prompt = reflection_prompt
-        self._initialize_agent()
+        self._action = self._build_graph()
 
     def generation_node(self, state: PlanningState) -> PlanningState:
         print("PlanningAgent: generating . . .")
@@ -48,7 +48,8 @@ class PlanningAgent(BaseAgent):
         return {
             "messages": [
                 self.llm.invoke(
-                    messages, {"configurable": {"thread_id": self.thread_id}}
+                    messages,
+                    self.build_config(tags=["planner", "generate"]),
                 )
             ]
         }
@@ -64,7 +65,8 @@ class PlanningAgent(BaseAgent):
         for _ in range(10):
             try:
                 res = self.llm.invoke(
-                    translated, {"configurable": {"thread_id": self.thread_id}}
+                    translated,
+                    self.build_config(tags=["planner", "formalize"]),
                 )
                 json_out = extract_json(res.content)
                 break
@@ -88,39 +90,72 @@ class PlanningAgent(BaseAgent):
         ]
         translated = [SystemMessage(content=reflection_prompt)] + translated
         res = self.llm.invoke(
-            translated, {"configurable": {"thread_id": self.thread_id}}
+            translated,
+            self.build_config(tags=["planner", "reflect"]),
         )
         return {"messages": [HumanMessage(content=res.content)]}
 
-    def _initialize_agent(self):
-        self.graph = StateGraph(PlanningState)
-        self.graph.add_node("generate", self.generation_node)
-        self.graph.add_node("reflect", self.reflection_node)
-        self.graph.add_node("formalize", self.formalize_node)
+    def _build_graph(self):
+        graph = StateGraph(PlanningState)
+        self.add_node(graph, self.generation_node, "generate")
+        self.add_node(graph, self.reflection_node, "reflect")
+        self.add_node(graph, self.formalize_node, "formalize")
 
-        self.graph.add_edge(START, "generate")
-        self.graph.add_edge("generate", "reflect")
-        self.graph.add_edge("formalize", END)
+        # Edges
+        graph.set_entry_point("generate")
+        graph.add_edge("generate", "reflect")
+        graph.set_finish_point("formalize")
 
-        self.graph.add_conditional_edges(
+        # Time the router logic too
+        graph.add_conditional_edges(
             "reflect",
-            should_continue,
+            self._wrap_cond(should_continue, "should_continue", "planner"),
             {"generate": "generate", "formalize": "formalize"},
         )
 
         # memory      = MemorySaver()
         # self.action = self.graph.compile(checkpointer=memory)
-        self.action = self.graph.compile(checkpointer=self.checkpointer)
+        return graph.compile(checkpointer=self.checkpointer)
         # self.action.get_graph().draw_mermaid_png(output_file_path="planning_agent_graph.png", draw_method=MermaidDrawMethod.PYPPETEER)
 
-    def run(self, prompt, recursion_limit=100):
-        initial_state = {"messages": [HumanMessage(content=prompt)]}
-        return self.action.invoke(
-            initial_state,
-            {
-                "recursion_limit": recursion_limit,
-                "configurable": {"thread_id": self.thread_id},
-            },
+    def _invoke(
+        self, inputs: Mapping[str, Any], recursion_limit: int = 1000, **_
+    ):
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
+        return self._action.invoke(inputs, config)
+
+    def _stream(
+        self,
+        inputs: Mapping[str, Any],
+        *,
+        config: dict | None = None,
+        recursion_limit: int = 1000,
+        **_,
+    ) -> Iterator[dict]:
+        # If you have defaults, merge them here:
+        default = self.build_config(
+            recursion_limit=recursion_limit, tags=["planner"]
+        )
+        if config:
+            merged = {**default, **config}
+            if "configurable" in config:
+                merged["configurable"] = {
+                    **default.get("configurable", {}),
+                    **config["configurable"],
+                }
+        else:
+            merged = default
+
+        # Delegate to the compiled graph's stream
+        yield from self._action.stream(inputs, merged)
+
+    # prevent bypass
+    @property
+    def action(self):
+        raise AttributeError(
+            "Use .stream(...) or .invoke(...); direct .action access is unsupported."
         )
 
 
@@ -137,15 +172,15 @@ def should_continue(state: PlanningState):
 
 def main():
     planning_agent = PlanningAgent()
-    for event in planning_agent.action.stream(
+
+    for event in planning_agent.stream(
         {
             "messages": [
                 HumanMessage(
-                    content="Find a city with as least 10 vowels in its name."  # "Write an essay on ideal high-entropy alloys for spacecraft."
+                    content="Find a city with at least 10 vowels in its name."
                 )
             ],
         },
-        config,
     ):
         print("-" * 30)
         print(event.keys())