ursa-ai 0.2.8__tar.gz → 0.2.9__tar.gz

{ursa_ai-0.2.8/src/ursa_ai.egg-info → ursa_ai-0.2.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ursa-ai
-Version: 0.2.8
+Version: 0.2.9
 Summary: Agents for science at LANL
 Author-email: Mike Grosskopf <mikegros@lanl.gov>, Nathan Debardeleben <ndebard@lanl.gov>, Rahul Somasundaram <rsomasundaram@lanl.gov>, Isaac Michaud <imichaud@lanl.gov>, Avanish Mishra <avanish@lanl.gov>, Arthur Lui <alui@lanl.gov>, Russell Bent <rbent@lanl.gov>, Earl Lawrence <earl@lanl.gov>
 License-Expression: BSD-3-Clause

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/pyproject.toml

@@ -65,11 +65,13 @@ enabled = true
 line-length = 80
 
 [tool.ruff.lint]
-ignore = ["D100"]
-extend-select = ["I", "W505"] # "D"
-extend-unsafe-fixes = ["F401"]
+extend-select = ["I"]  # W505 (doc-line-too-long); D (pydocstyle)
 pydocstyle.convention = "numpy"
 pycodestyle.max-doc-length = 80
+# extend-unfixable = ["F401"]  # unused imports
+
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
 
 # Ignore test file documentation linting.
 [tool.ruff.lint.extend-per-file-ignores]

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/__init__.py

@@ -1,9 +1,9 @@
-from .planning_agent import PlanningAgent, PlanningState
-from .websearch_agent import WebSearchAgent, WebSearchState
-from .execution_agent import ExecutionAgent, ExecutionState
+from .arxiv_agent import ArxivAgent, PaperMetadata, PaperState
+from .base import BaseAgent, BaseChatModel
 from .code_review_agent import CodeReviewAgent, CodeReviewState
+from .execution_agent import ExecutionAgent, ExecutionState
 from .hypothesizer_agent import HypothesizerAgent, HypothesizerState
-from .arxiv_agent import ArxivAgent, PaperState, PaperMetadata
-from .recall_agent import RecallAgent
-from .base import BaseAgent, BaseChatModel
 from .mp_agent import MaterialsProjectAgent
+from .planning_agent import PlanningAgent, PlanningState
+from .recall_agent import RecallAgent
+from .websearch_agent import WebSearchAgent, WebSearchState

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/arxiv_agent.py

@@ -1,29 +1,29 @@
+import base64
 import os
-import pymupdf
-import requests
-import feedparser
-from PIL import Image
+import re
+import statistics
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from io import BytesIO
-import base64
 from urllib.parse import quote
-from typing_extensions import TypedDict, List
-from concurrent.futures import ThreadPoolExecutor, as_completed
-from tqdm import tqdm
-import statistics
-import re
 
+import feedparser
+import pymupdf
+import requests
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_chroma import Chroma
 from langchain_community.document_loaders import PyPDFLoader
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from langgraph.graph import StateGraph, END, START
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_chroma import Chroma
+from langgraph.graph import StateGraph
+from PIL import Image
+from tqdm import tqdm
+from typing_extensions import List, TypedDict
 
 from .base import BaseAgent
 
 try:
     from openai import OpenAI
-except:
+except Exception:
     pass
 
 # embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
@@ -162,7 +162,7 @@ class ArxivAgent(BaseAgent):
                 full_id = entry.id.split("/abs/")[-1]
                 arxiv_id = full_id.split("/")[-1]
                 title = entry.title.strip()
-                authors = ", ".join(author.name for author in entry.authors)
+                # authors = ", ".join(author.name for author in entry.authors)
                 pdf_url = f"https://arxiv.org/pdf/{full_id}.pdf"
                 pdf_filename = os.path.join(
                     self.database_path, f"{arxiv_id}.pdf"
@@ -313,7 +313,7 @@ class ArxivAgent(BaseAgent):
 
         if "papers" not in state or len(state["papers"]) == 0:
             print(
-                f"No papers retrieved - bad query or network connection to ArXiv?"
+                "No papers retrieved - bad query or network connection to ArXiv?"
             )
             return {**state, "summaries": None}
 

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/base.py

@@ -1,9 +1,7 @@
 from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.load import dumps
 from langchain_litellm import ChatLiteLLM
 from langgraph.checkpoint.base import BaseCheckpointSaver
-from langchain_core.load import dumps
-
-import json
 
 
 class BaseAgent:

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/code_review_agent.py

@@ -7,7 +7,7 @@ from langchain_core.messages import HumanMessage, SystemMessage
 from langchain_core.tools import tool
 from langgraph.graph import END, START, StateGraph
 from langgraph.graph.message import add_messages
-from langgraph.prebuilt import ToolNode, InjectedState
+from langgraph.prebuilt import InjectedState, ToolNode
 from typing_extensions import TypedDict
 
 from ..prompt_library.code_review_prompts import (
@@ -349,7 +349,7 @@ def main():
     }
     result = (
         code_review_agent.action.invoke(initial_state),
-        {"configurable": {"thread_id": self.thread_id}},
+        {"configurable": {"thread_id": 42}},
     )
     for x in result["messages"]:
         print(x.content)

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/execution_agent.py

@@ -107,7 +107,7 @@ class ExecutionAgent(BaseAgent):
             # note that we've done the symlink now, so don't need to do it later
             new_state["symlinkdir"]["is_linked"] = True
 
-        if type(new_state["messages"][0]) == SystemMessage:
+        if isinstance(new_state["messages"][0], SystemMessage):
             new_state["messages"][0] = SystemMessage(
                 content=self.executor_prompt
             )
@@ -139,7 +139,7 @@ class ExecutionAgent(BaseAgent):
             memories = []
             # Handle looping through the messages
             for x in state["messages"]:
-                if not type(x) == AIMessage:
+                if not isinstance(x, AIMessage):
                     memories.append(x.content)
                 elif not x.tool_calls:
                     memories.append(x.content)
@@ -426,7 +426,7 @@ def edit_code(
 
     if old_code_clean not in content:
         console.print(
-            f"[yellow] ⚠️ 'old_code' not found in file'; no changes made.[/]"
+            "[yellow] ⚠️ 'old_code' not found in file'; no changes made.[/]"
         )
         return f"No changes made to {filename}: 'old_code' not found in file."
 
@@ -483,7 +483,7 @@ def command_safe(state: ExecutionState) -> Literal["safe", "unsafe"]:
     index = -1
     message = state["messages"][index]
     # Loop through all the consecutive tool messages in reverse order
-    while type(message) == ToolMessage:
+    while isinstance(message, ToolMessage):
         if "[UNSAFE]" in message.content:
             return "unsafe"
 

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/hypothesizer_agent.py

@@ -96,7 +96,7 @@ class HypothesizerAgent(BaseAgent):
             new_state["visited_sites"] = []
 
         try:
-            if type(raw_search_results) == str:
+            if isinstance(raw_search_results, str):
                 results_list = ast.literal_eval(raw_search_results)
             else:
                 results_list = raw_search_results
@@ -154,7 +154,7 @@ class HypothesizerAgent(BaseAgent):
             new_state["visited_sites"] = []
 
         try:
-            if type(raw_search_results) == str:
+            if isinstance(raw_search_results, str):
                 results_list = ast.literal_eval(raw_search_results)
             else:
                 results_list = raw_search_results
@@ -217,7 +217,7 @@ class HypothesizerAgent(BaseAgent):
             new_state["visited_sites"] = []
 
         try:
-            if type(raw_search_results) == str:
+            if isinstance(raw_search_results, str):
                 results_list = ast.literal_eval(raw_search_results)
             else:
                 results_list = raw_search_results
@@ -589,7 +589,7 @@ if __name__ == "__main__":
         initial_state,
         {
             "recursion_limit": 999999,
-            "configurable": {"thread_id": self.thread_id},
+            "configurable": {"thread_id": 42},
         },
     )
     summary_text = result["summary_report"]

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/mp_agent.py

@@ -1,20 +1,15 @@
-import os
 import json
-from typing import List, Dict
-from concurrent.futures import ThreadPoolExecutor
-from tqdm import tqdm
-
-from mp_api.client import MPRester
-
 import os
-from typing_extensions import TypedDict, List
-from concurrent.futures import ThreadPoolExecutor
-from tqdm import tqdm
 import re
+from concurrent.futures import ThreadPoolExecutor
+from typing import Dict
 
 from langchain_core.output_parsers import StrOutputParser
 from langchain_core.prompts import ChatPromptTemplate
-from langgraph.graph import StateGraph, END, START
+from langgraph.graph import StateGraph
+from mp_api.client import MPRester
+from tqdm import tqdm
+from typing_extensions import List, TypedDict
 
 from .base import BaseAgent
 

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/planning_agent.py

@@ -40,7 +40,7 @@ class PlanningAgent(BaseAgent):
 
     def generation_node(self, state: PlanningState) -> PlanningState:
         messages = state["messages"]
-        if type(messages[0]) == SystemMessage:
+        if isinstance(messages[0], SystemMessage):
             messages[0] = SystemMessage(content=self.planner_prompt)
         else:
             messages = [SystemMessage(content=self.planner_prompt)] + messages

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/recall_agent.py

@@ -1,4 +1,3 @@
-import os
 from .base import BaseAgent
 
 

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/agents/websearch_agent.py

@@ -1,5 +1,3 @@
-import inspect
-
 # from langchain_community.tools    import TavilySearchResults
 # from langchain_core.runnables.graph import MermaidDrawMethod
 from typing import Annotated, Any, List, Optional
@@ -9,7 +7,6 @@ from bs4 import BeautifulSoup
 from langchain_community.tools import DuckDuckGoSearchResults
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import HumanMessage, SystemMessage, ToolMessage
-from langchain_core.tools import tool
 from langchain_openai import ChatOpenAI
 from langgraph.graph import END, START, StateGraph
 from langgraph.graph.message import add_messages
@@ -19,8 +16,8 @@ from typing_extensions import TypedDict
 
 from ..prompt_library.websearch_prompts import (
     reflection_prompt,
-    websearch_prompt,
     summarize_prompt,
+    websearch_prompt,
 )
 from .base import BaseAgent
 
@@ -41,6 +38,7 @@ class WebSearchState(TypedDict):
     remaining_steps: int
     is_last_step: bool
     model: Any
+    thread_id: Any
 
 
 # Adding the model to the state clumsily so that all "read" sources arent in the
@@ -113,8 +111,14 @@ class WebSearchAgent(BaseAgent):
         except (requests.ConnectionError, requests.Timeout):
             return False
 
+    def state_store_node(self, state: WebSearchState) -> WebSearchState:
+        state["thread_id"] = self.thread_id
+        return state
+        # return dict(**state, thread_id=self.thread_id)
+
     def _initialize_agent(self):
         self.graph = StateGraph(WebSearchState)
+        self.graph.add_node("state_store", self.state_store_node)
         self.graph.add_node(
             "websearch",
             create_react_agent(
@@ -128,7 +132,8 @@ class WebSearchAgent(BaseAgent):
         self.graph.add_node("review", self.review_node)
         self.graph.add_node("response", self.response_node)
 
-        self.graph.add_edge(START, "websearch")
+        self.graph.add_edge(START, "state_store")
+        self.graph.add_edge("state_store", "websearch")
         self.graph.add_edge("websearch", "review")
         self.graph.add_edge("response", END)
 
@@ -185,7 +190,9 @@ def process_content(
     """
     summarized_information = (
         state["model"]
-        .invoke(content_prompt, {"configurable": {"thread_id": self.thread_id}})
+        .invoke(
+            content_prompt, {"configurable": {"thread_id": state["thread_id"]}}
+        )
         .content
     )
     return summarized_information
@@ -217,7 +224,7 @@ def main():
         inputs,
         {
             "recursion_limit": 10000,
-            "configurable": {"thread_id": self.thread_id},
+            "configurable": {"thread_id": 42},
         },
     )
 

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/prompt_library/hypothesizer_prompts.py

@@ -10,7 +10,7 @@ critic_prompt = dedent("""\
     You are Agent 2, a rigorous Critic who identifies flaws and areas for improvement.
 """)
 
-competitor_prompt = dedent(f"""\
+competitor_prompt = dedent("""\
     You are Agent 3, taking on the role of a direct competitor to Agent 1 in this hypothetical situation.
     Acting as that competitor, and taking into account potential critiques from the critic, provide an honest
     assessment how you might *REALLY* counter the approach of Agent 1.

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/util/memory_logger.py

@@ -1,7 +1,7 @@
 import os
 import shutil
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Sequence
+from typing import Any, Dict, Optional, Sequence
 
 from langchain_chroma import Chroma
 from langchain_core.documents import Document

{ursa_ai-0.2.8 → ursa_ai-0.2.9/src/ursa_ai.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ursa-ai
-Version: 0.2.8
+Version: 0.2.9
 Summary: Agents for science at LANL
 Author-email: Mike Grosskopf <mikegros@lanl.gov>, Nathan Debardeleben <ndebard@lanl.gov>, Rahul Somasundaram <rsomasundaram@lanl.gov>, Isaac Michaud <imichaud@lanl.gov>, Avanish Mishra <avanish@lanl.gov>, Arthur Lui <alui@lanl.gov>, Russell Bent <rbent@lanl.gov>, Earl Lawrence <earl@lanl.gov>
 License-Expression: BSD-3-Clause

{ursa_ai-0.2.8 → ursa_ai-0.2.9}/src/ursa/util/diff_renderer.py

@@ -1,5 +1,5 @@
-import re
 import difflib
+import re
 from dataclasses import dataclass
 
 from rich.console import Console, ConsoleOptions, RenderResult