fix examples

Author: luiarthur

.pre-commit-config.yaml

@@ -5,6 +5,8 @@ repos:
   hooks:
     # Run the linter.
     - id: ruff-check
-      args: [ --fix, --select, "I,F401" ]
+      args: [ --fix, --extend-fixable, "F401", --ignore, "D100,W505" ]
     # Run the formatter.
     - id: ruff-format
+
+# "src/**/__init__.py" = ["F401"]

examples/combine_agents_examples/agent_graph.py

@@ -54,14 +54,6 @@
 """
 
 
-class State(TypedDict):
-    messages: Annotated[list, add_messages]
-    current_progress: str
-    code_files: list[str]
-    workspace: str
-    arxiv_results: list[str]
-
-
 model = ChatLiteLLM(
     model="openai/o3",
     max_tokens=50000,
@@ -167,7 +159,7 @@ def runner(self, state: State) -> State:
             )
         os.makedirs(new_state["workspace"], exist_ok=True)
 
-        if type(new_state["messages"][0]) == SystemMessage:
+        if isinstance(new_state["messages"][0], SystemMessage):
             new_state["messages"][0] = SystemMessage(content=self.runner_prompt)
         else:
             new_state["messages"] = [
@@ -190,7 +182,7 @@ def summarize(self, state: ExecutionState) -> ExecutionState:
         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)

examples/single_agent_examples/arxiv_agent/high_entropy_alloy_papers.py

@@ -1,8 +1,5 @@
-import sys
 import time
 
-sys.path.append("../../.")
-
 from langchain_community.callbacks.openai_info import OpenAICallbackHandler
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
 
@@ -34,7 +31,7 @@ def main():
 
     t0 = time.time()
 
-    result = agent.run(
+    agent.run(
         arxiv_search_query="High Entropy Alloys",
         context="Find High entropy alloys suitable for application under extreme conditions. For candidates that you identify, provide the starting structure, crystal structure, lattice parameters, and space group.",
     )

examples/single_agent_examples/arxiv_agent/neutron_star_radius.py

@@ -1,3 +1,5 @@
+from pathlib import Path
+
 from langchain_litellm import ChatLiteLLM
 
 from ursa.agents import ArxivAgent
@@ -6,14 +8,15 @@
 def main():
     llm = ChatLiteLLM(model="openai/o3", max_completion_tokens=20000)
 
+    Path("workspace").mkdir(exist_ok=True)
     agent = ArxivAgent(
         llm=llm,
         summarize=True,
         process_images=True,
         max_results=3,
-        database_path="arxiv_papers_neutron_star",
-        summaries_path="arxiv_summaries_neutron_star",
-        vectorstore_path="arxiv_vectorstores_neutron_star",
+        database_path="workspace/arxiv_papers_neutron_star",
+        summaries_path="workspace/arxiv_summaries_neutron_star",
+        vectorstore_path="workspace/arxiv_vectorstores_neutron_star",
         download_papers=True,
     )
 

justfile

@@ -1,11 +1,14 @@
+ruff := "uvx ruff@0.12.10"
+lint-options := "--extend-fixable=F401 --ignore='D100,W505'"
+
 help:
     just -l -u
 
 test:
 	uv run examples/single_agent_examples/execution_agent/bayesian_optimization.py
 
 test-vowels:
-	uv run examples/single_agent_examples/research_agent/ten_vowel_city.py
+	uv run examples/single_agent_examples/websearch_agent/ten_vowel_city.py
 
 # Test neutron star example with latest dependencies
 neutron-latest:
@@ -19,5 +22,14 @@ clean-workspaces:
 	rm -rf workspace
 	rm -rf workspace_*/
 
-pre-commit:
+lint:
     uv run pre-commit run --all-files
+
+lint-diff:
+    {{ ruff }} check {{ lint-options }} --diff
+
+lint-stats:
+    {{ ruff }} check {{ lint-options }} --statistics
+
+lint-watch:
+    {{ ruff }} check {{ lint-options }} --watch

pyproject.toml

@@ -66,11 +66,14 @@ line-length = 80
 
 [tool.ruff.lint]
 ignore = ["D100"]
-extend-select = ["I", "W505"] # "D"
-extend-unsafe-fixes = ["F401"]
+extend-select = ["I", "W505"]  # D (docs)
+extend-unfixable = ["F401"]
 pydocstyle.convention = "numpy"
 pycodestyle.max-doc-length = 80
 
+[tool.ruff.lint.per-file-ignores]
+"__init__.py" = ["F401"]
+
 # Ignore test file documentation linting.
 [tool.ruff.lint.extend-per-file-ignores]
 "tests/**/*.py" = ["D"]

src/ursa/agents/__init__.py

@@ -1,29 +1,9 @@
-from ursa.agents.arxiv_agent import ArxivAgent, PaperMetadata, PaperState
-from ursa.agents.base import BaseAgent, BaseChatModel
-from ursa.agents.code_review_agent import CodeReviewAgent, CodeReviewState
-from ursa.agents.execution_agent import ExecutionAgent, ExecutionState
-from ursa.agents.hypothesizer_agent import HypothesizerAgent, HypothesizerState
-from ursa.agents.mp_agent import MaterialsProjectAgent
-from ursa.agents.planning_agent import PlanningAgent, PlanningState
-from ursa.agents.recall_agent import RecallAgent
-from ursa.agents.websearch_agent import WebSearchAgent, WebSearchState
-
-__all__ = [
-    "ArxivAgent",
-    "PaperMetadata",
-    "PaperState",
-    "BaseAgent",
-    "BaseChatModel",
-    "CodeReviewAgent",
-    "CodeReviewState",
-    "ExecutionAgent",
-    "ExecutionState",
-    "HypothesizerAgent",
-    "HypothesizerState",
-    "MaterialsProjectAgent",
-    "PlanningAgent",
-    "PlanningState",
-    "RecallAgent",
-    "WebSearchAgent",
-    "WebSearchState",
-]
+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 .mp_agent import MaterialsProjectAgent
+from .planning_agent import PlanningAgent, PlanningState
+from .recall_agent import RecallAgent
+from .websearch_agent import WebSearchAgent, WebSearchState

src/ursa/agents/arxiv_agent.py

@@ -23,7 +23,7 @@
 
 try:
     from openai import OpenAI
-except:
+except Exception:
     pass
 
 # embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
@@ -162,7 +162,7 @@ def _fetch_papers(self, query: str) -> List[PaperMetadata]:
                 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 @@ def process_paper(i, paper):
 
         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}
 

src/ursa/agents/code_review_agent.py

@@ -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)

src/ursa/agents/execution_agent.py

@@ -107,7 +107,7 @@ def query_executor(self, state: ExecutionState) -> ExecutionState:
             # 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 @@ def summarize(self, state: ExecutionState) -> ExecutionState:
             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"