ursa-ai 0.2.9__py3-none-any.whl → 0.2.11__py3-none-any.whl

ursa/agents/__init__.py

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

ursa/agents/arxiv_agent.py

@@ -211,12 +211,10 @@ class ArxivAgent(BaseAgent):
                 except Exception as e:
                     full_text = f"Error loading paper: {e}"
 
-            papers.append(
-                {
-                    "arxiv_id": arxiv_id,
-                    "full_text": full_text,
-                }
-            )
+            papers.append({
+                "arxiv_id": arxiv_id,
+                "full_text": full_text,
+            })
 
         return papers
 
@@ -279,28 +277,23 @@ class ArxivAgent(BaseAgent):
                     )
 
                     if relevant_docs_with_scores:
-                        score = sum(
-                            [s for _, s in relevant_docs_with_scores]
-                        ) / len(relevant_docs_with_scores)
+                        score = sum([
+                            s for _, s in relevant_docs_with_scores
+                        ]) / len(relevant_docs_with_scores)
                         relevancy_scores[i] = abs(1.0 - score)
                     else:
                         relevancy_scores[i] = 0.0
 
-                    retrieved_content = "\n\n".join(
-                        [
-                            doc.page_content
-                            for doc, _ in relevant_docs_with_scores
-                        ]
-                    )
+                    retrieved_content = "\n\n".join([
+                        doc.page_content for doc, _ in relevant_docs_with_scores
+                    ])
                 else:
                     retrieved_content = cleaned_text
 
-                summary = chain.invoke(
-                    {
-                        "retrieved_content": retrieved_content,
-                        "context": state["context"],
-                    }
-                )
+                summary = chain.invoke({
+                    "retrieved_content": retrieved_content,
+                    "context": state["context"],
+                })
 
             except Exception as e:
                 summary = f"Error summarizing paper: {e}"
@@ -376,9 +369,10 @@ 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"],
+        })
 
         with open(self.summaries_path + "/final_summary.txt", "w") as f:
             f.write(final_summary)
@@ -406,9 +400,10 @@ class ArxivAgent(BaseAgent):
         return graph
 
     def run(self, arxiv_search_query: str, context: str) -> str:
-        result = self.graph.invoke(
-            {"query": arxiv_search_query, "context": context}
-        )
+        result = self.graph.invoke({
+            "query": arxiv_search_query,
+            "context": context,
+        })
 
         if self.summarize:
             return result.get("final_summary", "No summary generated.")

ursa/agents/mp_agent.py

@@ -141,9 +141,10 @@ You are a materials-science assistant. Given the following metadata about a mate
         {context}
                 """)
         chain = prompt | self.llm | StrOutputParser()
-        final = chain.invoke(
-            {"summaries": combined, "context": state["context"]}
-        )
+        final = chain.invoke({
+            "summaries": combined,
+            "context": state["context"],
+        })
         return {**state, "final_summary": final}
 
     def _build_graph(self):

ursa/util/memory_logger.py

@@ -26,20 +26,30 @@ class AgentMemory:
     * Requires `langchain-chroma`, and `chromadb`.
     """
 
+    @classmethod
+    def get_db_path(cls, path: Optional[str | Path]) -> Path:
+        match path:
+            case None:
+                return Path.home() / ".cache" / "ursa" / "rag" / "db"
+            case str():
+                return Path(str)
+            case Path():
+                return path
+            case _:
+                raise TypeError(
+                    f"Type of path is `{type(path)}` "
+                    "but `Optional[str | Path]` was expected."
+                )
+
     def __init__(
         self,
         embedding_model,
         path: Optional[str | Path] = None,
         collection_name: str = "agent_memory",
     ) -> None:
-        self.path = (
-            Path(path)
-            if path
-            else Path(__file__).resolve().parent / "agent_memory_db"
-        )
+        self.path = self.get_db_path(path)
         self.collection_name = collection_name
         self.path.mkdir(parents=True, exist_ok=True)
-
         self.embeddings = embedding_model
 
         # If a DB already exists, load it; otherwise defer creation until `build_index`.
@@ -165,8 +175,7 @@ def delete_database(path: Optional[str | Path] = None):
         Where the on-disk Chroma DB is for deleting.  If *None*, a folder called
         ``agent_memory_db`` is created in the package’s base directory.
     """
-
-    db_path = Path(path) if path else Path("~/.cache/ursa/rag/db/")
+    db_path = AgentMemory.get_db_path(path)
     if os.path.exists(db_path):
         shutil.rmtree(db_path)
         print(f"Database: {db_path} has been deleted.")

{ursa_ai-0.2.9.dist-info → ursa_ai-0.2.11.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ursa-ai
-Version: 0.2.9
+Version: 0.2.11
 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.9.dist-info → ursa_ai-0.2.11.dist-info}/RECORD

@@ -1,10 +1,10 @@
-ursa/agents/__init__.py,sha256=PAPw6gjnmdA74HpiiYwtO5c6hYLgG_Yu22eu_nuGYKw,498
-ursa/agents/arxiv_agent.py,sha256=A_HVbmYOtQfalntOKSsUIkFu10oN9T6kBAvCuAz2evg,14273
+ursa/agents/__init__.py,sha256=IdwVBLlz7QydHiFSok06McIfEJed6rxrT01hZgpQig0,1008
+ursa/agents/arxiv_agent.py,sha256=6pycpm1hkUNjwXF6lJ4DCF5CxOiSnS5dvpW2Qes3ZD8,14125
 ursa/agents/base.py,sha256=uFhRLVzqhFbTZVA7IePKbUi03ATCXuvga7rzwaHy1B0,1321
 ursa/agents/code_review_agent.py,sha256=aUDq5gT-jdl9Qs-Wewj2oz1d60xov9sN-DOYRfGNTU0,11550
 ursa/agents/execution_agent.py,sha256=-At1EcKRHP9lYQ80jpqTPtQyPuQV-sIpf9J8LlEfWdA,16618
 ursa/agents/hypothesizer_agent.py,sha256=rSLohNQz3xvEcL_DGTFivf9q5BlX1cqlLUcts4GJIjM,23309
-ursa/agents/mp_agent.py,sha256=HTMAnv1yGs8vgRLGFFYHSbwOz24qdnB-if_JQSH3urQ,6002
+ursa/agents/mp_agent.py,sha256=UyJSheMGHZpWQJL3EgYgPPqArfv6F8sndN05q4CPtyo,6015
 ursa/agents/planning_agent.py,sha256=AKWQJ848RLPiwQGrvDNdN9lBlf3YI5qWmt2hqXnRGj8,5426
 ursa/agents/recall_agent.py,sha256=bQk7ZJtiO5pj89A50OBDzAJ4G2F7ZdsMwmKnp1WWR7g,813
 ursa/agents/websearch_agent.py,sha256=rCv4AWbqe5Us4FmuypM6jptri21nKoNg044ncsu9u3E,8014
@@ -17,10 +17,10 @@ ursa/prompt_library/websearch_prompts.py,sha256=n4DJaYn_lIYAVtdy00CCJjT-dLWhn2JN
 ursa/tools/run_command.py,sha256=sQRuHtRyJYWEyL9dpW_Ukc-xQ5vmKKJK1i_6z3uKEfA,690
 ursa/tools/write_code.py,sha256=DtCsUMZegYm0mk-HMPG5Zo3Ba1gbGfnXHsv1NZTdDs8,1220
 ursa/util/diff_renderer.py,sha256=1L1q2qWWb8gLhR532-LgJn2TrqXDx0gUpPVOWD_sqeU,4086
-ursa/util/memory_logger.py,sha256=-4jZkMFXLnABj9x_DMGEUySVPLZaI47HrLgK69Naxw0,5731
+ursa/util/memory_logger.py,sha256=9tNPF3Fgth4ImpNbs7Qg5Ci6Hb1ziMWsDTsqjkDi1iY,6096
 ursa/util/parse.py,sha256=M0cjyQWmjatxX4WbVmDRUiirTLyW-t_Aemlrlrsc5nA,2811
-ursa_ai-0.2.9.dist-info/licenses/LICENSE,sha256=4Vr6_u2zTHIUvYjoOBg9ztDbfpV3hyCFv3mTCS87gYU,1482
-ursa_ai-0.2.9.dist-info/METADATA,sha256=tPMX8nDJxo5ZbaOnpScf9wQrkTsUD73XrStyHxs81Eg,6848
-ursa_ai-0.2.9.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ursa_ai-0.2.9.dist-info/top_level.txt,sha256=OjA1gRYSUAeiXGnpqPC8iOOGfcjFO1IlP848qMnYSdY,5
-ursa_ai-0.2.9.dist-info/RECORD,,
+ursa_ai-0.2.11.dist-info/licenses/LICENSE,sha256=4Vr6_u2zTHIUvYjoOBg9ztDbfpV3hyCFv3mTCS87gYU,1482
+ursa_ai-0.2.11.dist-info/METADATA,sha256=apeOKu1ZbVz8Fbi4p41DsdGLkasx8yZ0uuE2BjUAmek,6849
+ursa_ai-0.2.11.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+ursa_ai-0.2.11.dist-info/top_level.txt,sha256=OjA1gRYSUAeiXGnpqPC8iOOGfcjFO1IlP848qMnYSdY,5
+ursa_ai-0.2.11.dist-info/RECORD,,