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

ursa/agents/arxiv_agent.py

@@ -3,12 +3,14 @@ import os
 import re
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from io import BytesIO
+from typing import Any, Mapping
 from urllib.parse import quote
 
 import feedparser
 import pymupdf
 import requests
 from langchain_community.document_loaders import PyPDFLoader
+from langchain_core.language_models import BaseChatModel
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
 from langgraph.graph import StateGraph
@@ -16,8 +18,8 @@ from PIL import Image
 from tqdm import tqdm
 from typing_extensions import List, TypedDict
 
-from .base import BaseAgent
-from .rag_agent import RAGAgent
+from ursa.agents.base import BaseAgent
+from ursa.agents.rag_agent import RAGAgent
 
 try:
     from openai import OpenAI
@@ -120,7 +122,7 @@ def remove_surrogates(text: str) -> str:
 class ArxivAgent(BaseAgent):
     def __init__(
         self,
-        llm="openai/o3-mini",
+        llm: str | BaseChatModel = "openai/o3-mini",
         summarize: bool = True,
         process_images=True,
         max_results: int = 3,
@@ -141,7 +143,7 @@ class ArxivAgent(BaseAgent):
         self.download_papers = download_papers
         self.rag_embedding = rag_embedding
 
-        self.graph = self._build_graph()
+        self._action = self._build_graph()
 
         os.makedirs(self.database_path, exist_ok=True)
 
@@ -259,10 +261,13 @@ class ArxivAgent(BaseAgent):
 
             try:
                 cleaned_text = remove_surrogates(paper["full_text"])
-                summary = chain.invoke({
-                    "retrieved_content": cleaned_text,
-                    "context": state["context"],
-                })
+                summary = chain.invoke(
+                    {
+                        "retrieved_content": cleaned_text,
+                        "context": state["context"],
+                    },
+                    config=self.build_config(tags=["arxiv", "summarize_each"]),
+                )
 
             except Exception as e:
                 summary = f"Error summarizing paper: {e}"
@@ -304,7 +309,9 @@ class ArxivAgent(BaseAgent):
             embedding=self.rag_embedding,
             database_path=self.database_path,
         )
-        new_state["final_summary"] = rag_agent.run(context=state["context"])
+        new_state["final_summary"] = rag_agent.invoke(context=state["context"])[
+            "summary"
+        ]
         return new_state
 
     def _aggregate_node(self, state: PaperState) -> PaperState:
@@ -341,10 +348,13 @@ class ArxivAgent(BaseAgent):
 
         chain = prompt | self.llm | StrOutputParser()
 
-        final_summary = chain.invoke({
-            "Summaries": combined,
-            "context": state["context"],
-        })
+        final_summary = chain.invoke(
+            {
+                "Summaries": combined,
+                "context": state["context"],
+            },
+            config=self.build_config(tags=["arxiv", "aggregate"]),
+        )
 
         with open(self.summaries_path + "/final_summary.txt", "w") as f:
             f.write(final_summary)
@@ -352,49 +362,69 @@ class ArxivAgent(BaseAgent):
         return {**state, "final_summary": final_summary}
 
     def _build_graph(self):
-        builder = StateGraph(PaperState)
-        builder.add_node("fetch_papers", self._fetch_node)
+        graph = StateGraph(PaperState)
 
+        self.add_node(graph, self._fetch_node)
         if self.summarize:
             if self.rag_embedding:
-                builder.add_node("rag_summarize", self._rag_node)
-
-                builder.set_entry_point("fetch_papers")
-                builder.add_edge("fetch_papers", "rag_summarize")
-                builder.set_finish_point("rag_summarize")
+                self.add_node(graph, self._rag_node)
+                graph.set_entry_point("_fetch_node")
+                graph.add_edge("_fetch_node", "_rag_node")
+                graph.set_finish_point("_rag_node")
             else:
-                builder.add_node("summarize_each", self._summarize_node)
-                builder.add_node("aggregate", self._aggregate_node)
-
-                builder.set_entry_point("fetch_papers")
-                builder.add_edge("fetch_papers", "summarize_each")
-                builder.add_edge("summarize_each", "aggregate")
-                builder.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:
-            builder.set_entry_point("fetch_papers")
-            builder.set_finish_point("fetch_papers")
+            graph.set_entry_point("_fetch_node")
+            graph.set_finish_point("_fetch_node")
 
-        graph = builder.compile()
-        return graph
+        return graph.compile(checkpointer=self.checkpointer)
 
-    def run(self, arxiv_search_query: str, context: str) -> str:
-        result = self.graph.invoke({
-            "query": arxiv_search_query,
-            "context": context,
-        })
+    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 self.summarize:
-            return result.get("final_summary", "No summary generated.")
-        else:
-            return "\n\nFinished Fetching papers!"
+        # this seems dumb, but it's b/c sometimes we had referred to the value as
+        # 'query' other times as 'arxiv_search_query' so trying to keep it compatible
+        # aliasing: accept arxiv_search_query -> query
+        if "query" not in inputs:
+            if "arxiv_search_query" in inputs:
+                # make a shallow copy and rename the key
+                inputs = dict(inputs)
+                inputs["query"] = inputs.pop("arxiv_search_query")
+            else:
+                raise KeyError(
+                    "Missing 'query' in inputs (alias 'arxiv_search_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 papers!"
+        )
 
-if __name__ == "__main__":
-    agent = ArxivAgent()
-    result = agent.run(
-        arxiv_search_query="Experimental Constraints on neutron star radius",
-        context="What are the constraints on the neutron star radius and what uncertainties are there on the constraints?",
-    )
 
-    print(result)
+# NOTE: Run test in `tests/agents/test_arxiv_agent/test_arxiv_agent.py` via:
+#
+# pytest -s tests/agents/test_arxiv_agent
+#
+# OR
+#
+# uv run pytest -s tests/agents/test_arxiv_agent

ursa/agents/base.py

@@ -1,10 +1,47 @@
+import re
+from abc import ABC, abstractmethod
+from contextvars import ContextVar
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    Mapping,
+    Optional,
+    Sequence,
+    Union,
+    final,
+)
+from uuid import uuid4
+
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.load import dumps
+from langchain_core.runnables import (
+    RunnableLambda,
+)
 from langchain_litellm import ChatLiteLLM
 from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.graph import StateGraph
+
+from ursa.observability.timing import (
+    Telemetry,  # for timing / telemetry / metrics
+)
+
+InputLike = Union[str, Mapping[str, Any]]
+_INVOKE_DEPTH = ContextVar("_INVOKE_DEPTH", default=0)
 
 
-class BaseAgent:
+def _to_snake(s: str) -> str:
+    s = re.sub(
+        r"^([A-Z]{2,})([A-Z][a-z])",
+        lambda m: m.group(1)[0] + m.group(1)[1:].lower() + m.group(2),
+        str(s),
+    )  # RAGAgent -> RagAgent
+    s = re.sub(r"(?<!^)(?=[A-Z])", "_", s)  # CamelCase -> snake_case
+    s = s.replace("-", "_").replace(" ", "_")
+    return s.lower()
+
+
+class BaseAgent(ABC):
     # llm: BaseChatModel
     # llm_with_tools: Runnable[LanguageModelInput, BaseMessage]
 
@@ -12,6 +49,10 @@ class BaseAgent:
         self,
         llm: str | BaseChatModel,
         checkpointer: BaseCheckpointSaver = None,
+        enable_metrics: bool = False,  # default to enabling metrics
+        metrics_dir: str = ".ursa_metrics",  # dir to save metrics, with a default
+        autosave_metrics: bool = True,
+        thread_id: Optional[str] = None,
         **kwargs,
     ):
         match llm:
@@ -32,10 +73,336 @@ class BaseAgent:
                     "llm argument must be a string with the provider and model, or a BaseChatModel instance."
                 )
 
+        self.thread_id = thread_id or uuid4().hex
         self.checkpointer = checkpointer
-        self.thread_id = self.__class__.__name__
+        self.telemetry = Telemetry(
+            enable=enable_metrics,
+            output_dir=metrics_dir,
+            save_json_default=autosave_metrics,
+        )
+
+    @property
+    def name(self) -> str:
+        """Agent name."""
+        return self.__class__.__name__
+
+    def add_node(
+        self,
+        graph: StateGraph,
+        f: Callable[..., Mapping[str, Any]],
+        node_name: Optional[str] = None,
+        agent_name: Optional[str] = None,
+    ) -> StateGraph:
+        """Add node to graph.
+
+        This is used to track token usage and is simply the following.
+
+        ```python
+        _node_name = node_name or f.__name__
+        return graph.add_node(
+            _node_name, self._wrap_node(f, _node_name, self.name)
+        )
+        ```
+        """
+        _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 graph.add_node(_node_name, wrapped_node)
 
     def write_state(self, filename, state):
         json_state = dumps(state, ensure_ascii=False)
         with open(filename, "w") as f:
             f.write(json_state)
+
+    # BaseAgent
+    def build_config(self, **overrides) -> dict:
+        """
+        Build a config dict that includes telemetry callbacks and the thread_id.
+        You can pass overrides like recursion_limit=..., configurable={...}, etc.
+        """
+        base = {
+            "configurable": {"thread_id": self.thread_id},
+            "metadata": {
+                "thread_id": self.thread_id,
+                "telemetry_run_id": self.telemetry.context.get("run_id"),
+            },
+            # "configurable": {
+            #     "thread_id": getattr(self, "thread_id", "default")
+            # },
+            # "metadata": {
+            #     "thread_id": getattr(self, "thread_id", "default"),
+            #     "telemetry_run_id": self.telemetry.context.get("run_id"),
+            # },
+            "tags": [self.name],
+            "callbacks": self.telemetry.callbacks,
+        }
+        # include model name when we can
+        model_name = getattr(self, "llm_model", None) or getattr(
+            getattr(self, "llm", None), "model", None
+        )
+        if model_name:
+            base["metadata"]["model"] = model_name
+
+        if "configurable" in overrides and isinstance(
+            overrides["configurable"], dict
+        ):
+            base["configurable"].update(overrides.pop("configurable"))
+        if "metadata" in overrides and isinstance(overrides["metadata"], dict):
+            base["metadata"].update(overrides.pop("metadata"))
+        # merge tags if caller provides them
+        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"]
+            ]
+        base.update(overrides)
+        return base
+
+    # agents will invoke like this:
+    # planning_output = planner.invoke(
+    #     {"messages": [HumanMessage(content=problem)]},
+    #     config={
+    #         "recursion_limit": 999_999,
+    #         "configurable": {"thread_id": planner.thread_id},
+    #     },
+    # )
+    # they can also, separately, override these defaults about metrics
+    # keys that are NOT inputs; they should not be folded into the inputs mapping
+    _TELEMETRY_KW = {
+        "raw_debug",
+        "save_json",
+        "metrics_path",
+        "save_raw_snapshot",
+        "save_raw_records",
+    }
+    _CONTROL_KW = {"config", "recursion_limit", "tags", "metadata", "callbacks"}
+
+    @final
+    def invoke(
+        self,
+        inputs: Optional[InputLike] = None,  # sentinel
+        /,
+        *,
+        raw_debug: bool = False,
+        save_json: 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,  # may contain inputs (keyword-inputs) and/or control kw
+    ) -> Any:
+        depth = _INVOKE_DEPTH.get()
+        _INVOKE_DEPTH.set(depth + 1)
+        try:
+            if depth == 0:
+                self.telemetry.begin_run(
+                    agent=self.name, thread_id=self.thread_id
+                )
+
+            # If no positional inputs were provided, split kwargs into inputs vs control
+            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  # only control kwargs remain
+
+            # If both positional inputs and extra unknown kwargs-as-inputs are given, forbid merging
+            else:
+                # keep only control kwargs; anything else would be ambiguous
+                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."
+                        )
+
+            # subclasses may translate keys
+            normalized = self._normalize_inputs(inputs)
+
+            # forward config + any control kwargs (e.g., recursion_limit) to the agent
+            return self._invoke(normalized, config=config, **kwargs)
+
+        finally:
+            new_depth = _INVOKE_DEPTH.get() - 1
+            _INVOKE_DEPTH.set(new_depth)
+            if new_depth == 0:
+                self.telemetry.render(
+                    raw=raw_debug,
+                    save_json=save_json,
+                    filepath=metrics_path,
+                    save_raw_snapshot=save_raw_snapshot,
+                    save_raw_records=save_raw_records,
+                )
+
+    def _normalize_inputs(self, inputs: InputLike) -> Mapping[str, Any]:
+        if isinstance(inputs, str):
+            # Adjust to your message type
+            from langchain_core.messages import HumanMessage
+
+            return {"messages": [HumanMessage(content=inputs)]}
+        if isinstance(inputs, Mapping):
+            return inputs
+        raise TypeError(f"Unsupported input type: {type(inputs)}")
+
+    @abstractmethod
+    def _invoke(self, inputs: Mapping[str, Any], **config: Any) -> Any:
+        """Subclasses implement the actual work against normalized inputs."""
+        ...
+
+    def __call__(self, inputs: InputLike, /, **kwargs: Any) -> Any:
+        return self.invoke(inputs, **kwargs)
+
+    # Runtime enforcement: forbid subclasses from overriding invoke
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        if "invoke" in cls.__dict__:
+            raise TypeError(
+                f"{cls.__name__} must not override BaseAgent.invoke(); implement _invoke() only."
+            )
+
+    def stream(
+        self,
+        inputs: InputLike,
+        config: Any | None = None,  # allow positional/keyword like LangGraph
+        /,
+        *,
+        raw_debug: bool = False,
+        save_json: bool | None = None,
+        metrics_path: str | None = None,
+        save_raw_snapshot: bool | None = None,
+        save_raw_records: bool | None = None,
+        **kwargs: Any,
+    ) -> Iterator[Any]:
+        """Public streaming entry point. Telemetry-wrapped."""
+        depth = _INVOKE_DEPTH.get()
+        _INVOKE_DEPTH.set(depth + 1)
+        try:
+            if depth == 0:
+                self.telemetry.begin_run(
+                    agent=self.name, thread_id=self.thread_id
+                )
+            normalized = self._normalize_inputs(inputs)
+            yield from self._stream(normalized, config=config, **kwargs)
+        finally:
+            new_depth = _INVOKE_DEPTH.get() - 1
+            _INVOKE_DEPTH.set(new_depth)
+            if new_depth == 0:
+                self.telemetry.render(
+                    raw=raw_debug,
+                    save_json=save_json,
+                    filepath=metrics_path,
+                    save_raw_snapshot=save_raw_snapshot,
+                    save_raw_records=save_raw_records,
+                )
+
+    def _stream(
+        self,
+        inputs: Mapping[str, Any],
+        *,
+        config: Any | None = None,
+        **kwargs: Any,
+    ) -> Iterator[Any]:
+        raise NotImplementedError(
+            f"{self.name} does not support streaming. "
+            "Override _stream(...) in your agent to enable it."
+        )
+
+    # def run(
+    #     self,
+    #     *args,
+    #     raw_debug: bool = False,
+    #     save_json: bool | None = None,
+    #     metrics_path: str | None = None,
+    #     save_raw_snapshot: bool | None = None,
+    #     save_raw_records: bool | None = None,
+    #     **kwargs
+    # ):
+    #     try:
+    #         self.telemetry.begin_run(agent=self.name, thread_id=self.thread_id)
+    #         result = self._run_impl(*args, **kwargs)
+    #         return result
+    #     finally:
+    #         print(self.telemetry.render(
+    #             raw=raw_debug,
+    #             save_json=save_json,
+    #             filepath=metrics_path,
+    #             save_raw_snapshot=save_raw_snapshot,
+    #             save_raw_records=save_raw_records,
+    #         ))
+
+    # @abstractmethod
+    # def _run_impl(self, *args, **kwargs):
+    #     raise NotImplementedError("Agents must implement _run_impl")
+
+    def _default_node_tags(
+        self, name: str, extra: Sequence[str] | None = None
+    ) -> list[str]:
+        tags = [self.name, "graph", name]
+        if extra:
+            tags.extend(extra)
+        return tags
+
+    def _as_runnable(self, fn: Any):
+        # If it's already runnable (has .with_config/.invoke), return it; else wrap
+        return (
+            fn
+            if hasattr(fn, "with_config") and hasattr(fn, "invoke")
+            else RunnableLambda(fn)
+        )
+
+    def _node_cfg(self, name: str, *extra_tags: str) -> dict:
+        """Build a consistent config for a node/runnable so we can reapply it after .map(), subgraph compile, etc."""
+        ns = extra_tags[0] if extra_tags else _to_snake(self.name)
+        tags = [self.name, "graph", name, *extra_tags]
+        return dict(
+            run_name="node",  # keep "node:" prefixing in the timer; don't fight Rich labels here
+            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 our node config applied. Safe to call on callables or runnables.
+        IMPORTANT: call this AGAIN after .map() / subgraph .compile() (they often drop config)."""
+        r = self._as_runnable(runnable_or_fn)
+        return r.with_config(**self._node_cfg(name, *extra_tags))
+
+    def _wrap_node(self, fn_or_runnable, name: str, *extra_tags: str):
+        return self.ns(fn_or_runnable, name, *extra_tags)
+
+    def _wrap_cond(self, fn: Any, name: str, *extra_tags: str):
+        ns = extra_tags[0] if extra_tags else _to_snake(self.name)
+        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):
+        ns = extra_tags[0] if extra_tags else _to_snake(self.name)
+        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,
+            },
+        )