ursa-ai 0.0.3__py3-none-any.whl → 0.2.2__py3-none-any.whl

ursa/agents/mp_agent.py

@@ -0,0 +1,257 @@
+import os
+import json
+from typing import List, Dict
+from concurrent.futures import ThreadPoolExecutor
+from tqdm import tqdm
+
+from mp_api.client import MPRester
+from langchain.schema import Document
+
+import os
+import pymupdf  
+import requests
+import feedparser
+from PIL import Image
+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 re
+
+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 langchain_openai import OpenAIEmbeddings
+
+from openai import OpenAI
+
+from .base import BaseAgent
+
+
+
+
+client = OpenAI()
+
+embeddings = OpenAIEmbeddings()
+
+class PaperMetadata(TypedDict):
+    arxiv_id: str
+    full_text: str
+
+class PaperState(TypedDict, total=False):
+    query: str
+    context: str
+    papers: List[PaperMetadata]
+    summaries: List[str]
+    final_summary: str
+
+
+def describe_image(image: Image.Image) -> str:
+    buffered = BytesIO()
+    image.save(buffered, format="PNG")
+    img_base64 = base64.b64encode(buffered.getvalue()).decode()
+
+    response = client.chat.completions.create(
+        model="gpt-4-vision-preview",
+        messages=[
+            {"role": "system", "content": "You are a scientific assistant who explains plots and scientific diagrams."},
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": "Describe this scientific image or plot in detail."},
+                    {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{img_base64}"}}
+                ],
+            },
+        ],
+        max_tokens=500,
+    )
+    return response.choices[0].message.content.strip()
+
+
+def extract_and_describe_images(pdf_path: str, max_images: int = 5) -> List[str]:
+    doc = pymupdf.open(pdf_path)
+    descriptions = []
+    image_count = 0
+
+    for page_index in range(len(doc)):
+        if image_count >= max_images:
+            break
+        page = doc[page_index]
+        images = page.get_images(full=True)
+
+        for img_index, img in enumerate(images):
+            if image_count >= max_images:
+                break
+            xref = img[0]
+            base_image = doc.extract_image(xref)
+            image_bytes = base_image["image"]
+            image = Image.open(BytesIO(image_bytes))
+
+            try:
+                desc = describe_image(image)
+                descriptions.append(f"Page {page_index + 1}, Image {img_index + 1}: {desc}")
+            except Exception as e:
+                descriptions.append(f"Page {page_index + 1}, Image {img_index + 1}: [Error: {e}]")
+            image_count += 1
+
+    return descriptions
+
+
+def remove_surrogates(text: str) -> str:
+    return re.sub(r'[\ud800-\udfff]', '', text)
+
+class MaterialsProjectAgent(BaseAgent):
+    def __init__(
+        self,
+        llm="openai/o3-mini",
+        summarize: bool = True,
+        max_results: int = 3,
+        database_path: str = 'mp_database',
+        summaries_path: str = 'mp_summaries',
+        vectorstore_path: str = 'mp_vectorstores',
+        **kwargs
+    ):
+        super().__init__(llm, **kwargs)
+        self.summarize        = summarize
+        self.max_results      = max_results
+        self.database_path    = database_path
+        self.summaries_path   = summaries_path
+        self.vectorstore_path = vectorstore_path
+        
+        os.makedirs(self.database_path,    exist_ok=True)
+        os.makedirs(self.summaries_path,   exist_ok=True)
+        os.makedirs(self.vectorstore_path, exist_ok=True)
+
+        self.embeddings = OpenAIEmbeddings()  # or your preferred embedding
+        self.graph = self._build_graph()
+
+    def _fetch_node(self, state: Dict) -> Dict:
+        f   = state["query"]
+        els = f["elements"]                  # e.g. ["Ga","In"]
+        bg  = (f["band_gap_min"], f["band_gap_max"])
+        e_above_hull = (0, 0)                # only on-hull (stable)
+        mats = []
+        with MPRester() as mpr:
+            # get ALL matching materials…
+            all_results = mpr.materials.summary.search(
+                elements=els,
+                band_gap=bg,
+                energy_above_hull=e_above_hull,
+                is_stable=True              # equivalent filter
+            )
+            # …then take only the first `max_results`
+            for doc in all_results[: self.max_results]:
+                mid = doc.material_id
+                data = doc.dict()
+                # cache to disk
+                path = os.path.join(self.database_path, f"{mid}.json")
+                if not os.path.exists(path):
+                    with open(path, "w") as f:
+                        json.dump(data, f, indent=2)
+                mats.append({"material_id": mid, "metadata": data})
+
+        return {**state, "materials": mats}
+
+    def _get_or_build_vectorstore(self, text: str, mid: str):
+        """Build or load a Chroma vectorstore for a single material's description."""
+        persist_dir = os.path.join(self.vectorstore_path, mid)
+        if os.path.exists(persist_dir):
+            store = Chroma(persist_directory=persist_dir, embedding_function=self.embeddings)
+        else:
+            splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
+            docs = splitter.create_documents([text])
+            store = Chroma.from_documents(docs, self.embeddings, persist_directory=persist_dir)
+        return store.as_retriever(search_kwargs={"k": 5})
+
+    def _summarize_node(self, state: Dict) -> Dict:
+        """Summarize each material via LLM over its metadata."""
+        # prompt template
+        prompt = ChatPromptTemplate.from_template("""
+You are a materials-science assistant. Given the following metadata about a material, produce a concise summary focusing on its key properties:
+
+{metadata}
+        """)
+        chain = prompt | self.llm | StrOutputParser()
+
+        summaries = [None] * len(state["materials"])
+        relevancy = [0.0] * len(state["materials"])
+
+        def process(i, mat):
+            mid = mat["material_id"]
+            meta = mat["metadata"]
+            # flatten metadata to text
+            text = "\n".join(f"{k}: {v}" for k, v in meta.items())
+            # build or load summary
+            summary_file = os.path.join(self.summaries_path, f"{mid}_summary.txt")
+            if os.path.exists(summary_file):
+                with open(summary_file) as f:
+                    return i, f.read()
+            # optional: vectorize & retrieve, but here we just summarize full text
+            result = chain.invoke({"metadata": text})
+            with open(summary_file, 'w') as f:
+                f.write(result)
+            return i, result
+
+        with ThreadPoolExecutor(max_workers=min(8, len(state["materials"]))) as exe:
+            futures = [exe.submit(process, i, m) for i, m in enumerate(state["materials"])]
+            for future in tqdm(futures, desc="Summarizing materials"):
+                i, summ = future.result()
+                summaries[i] = summ
+
+        return {**state, "summaries": summaries}
+
+    def _aggregate_node(self, state: Dict) -> Dict:
+        """Combine all summaries into a single, coherent answer."""
+        combined = "\n\n----\n\n".join(
+            f"[{i+1}] {m['material_id']}\n\n{summary}"
+            for i, (m, summary) in enumerate(zip(state["materials"], state["summaries"]))
+        )
+
+        prompt = ChatPromptTemplate.from_template("""
+        You are a materials informatics assistant. Below are brief summaries of several materials:
+
+        {summaries}
+
+        Answer the user’s question in context:
+
+        {context}
+                """)
+        chain = prompt | self.llm | StrOutputParser()
+        final = chain.invoke({"summaries": combined, "context": state["context"]})
+        return {**state, "final_summary": final}
+
+    def _build_graph(self):
+        g = StateGraph(dict)  # using plain dict for state
+        g.add_node("fetch", self._fetch_node)
+        if self.summarize:
+            g.add_node("summarize", self._summarize_node)
+            g.add_node("aggregate", self._aggregate_node)
+            g.set_entry_point("fetch")
+            g.add_edge("fetch", "summarize")
+            g.add_edge("summarize", "aggregate")
+            g.set_finish_point("aggregate")
+        else:
+            g.set_entry_point("fetch")
+            g.set_finish_point("fetch")
+        return g.compile()
+
+    def run(self, mp_query: str, context: str) -> str:
+        state = {"query": mp_query, "context": context}
+        out = self.graph.invoke(state)
+        if self.summarize:
+            return out.get("final_summary", "")
+        return json.dumps(out.get("materials", []), indent=2)
+
+
+if __name__ == "__main__":
+    agent = MaterialsProjectAgent()
+    resp = agent.run(
+        mp_query="LiFePO4",
+        context="What is its band gap and stability, and any synthesis challenges?"
+    )
+    print(resp)

ursa/agents/planning_agent.py

@@ -0,0 +1,138 @@
+# from langgraph.checkpoint.memory  import MemorySaver
+# from langchain_core.runnables.graph import MermaidDrawMethod
+from typing import Annotated, Any, Dict, List, 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.message import add_messages
+from pydantic import Field
+from typing_extensions import TypedDict
+
+from ..prompt_library.planning_prompts import (
+    formalize_prompt,
+    planner_prompt,
+    reflection_prompt,
+)
+from ..util.parse import extract_json
+from .base import BaseAgent
+
+
+class PlanningState(TypedDict):
+    messages: Annotated[list, add_messages]
+    plan_steps: List[Dict[str, Any]] = Field(
+        default_factory=list, description="Ordered steps in the solution plan"
+    )
+    reflection_steps: Optional[int] = Field(
+        default=3, description="Number of reflection steps"
+    )
+
+
+class PlanningAgent(BaseAgent):
+    def __init__(
+        self, llm: str | BaseChatModel = "openai/gpt-4o-mini", **kwargs
+    ):
+        super().__init__(llm, **kwargs)
+        self.planner_prompt = planner_prompt
+        self.formalize_prompt = formalize_prompt
+        self.reflection_prompt = reflection_prompt
+        self._initialize_agent()
+
+    def generation_node(self, state: PlanningState) -> PlanningState:
+        messages = state["messages"]
+        if type(messages[0]) == SystemMessage:
+            messages[0] = SystemMessage(content=self.planner_prompt)
+        else:
+            messages = [SystemMessage(content=self.planner_prompt)] + messages
+        return {"messages": [self.llm.invoke(messages, {"configurable": {"thread_id": self.thread_id}})]}
+
+    def formalize_node(self, state: PlanningState) -> PlanningState:
+        cls_map = {"ai": HumanMessage, "human": AIMessage}
+        translated = [state["messages"][0]] + [
+            cls_map[msg.type](content=msg.content)
+            for msg in state["messages"][1:]
+        ]
+        translated = [SystemMessage(content=self.formalize_prompt)] + translated
+        for _ in range(10):
+            try:
+                res = self.llm.invoke(translated, {"configurable": {"thread_id": self.thread_id}})
+                json_out = extract_json(res.content)
+                break
+            except ValueError:
+                translated.append(
+                    HumanMessage(
+                        content="Your response was not valid JSON. Try again."
+                    )
+                )
+        return {
+            "messages": [HumanMessage(content=res.content)],
+            "plan_steps": json_out,
+        }
+
+    def reflection_node(self, state: PlanningState) -> PlanningState:
+        cls_map = {"ai": HumanMessage, "human": AIMessage}
+        translated = [state["messages"][0]] + [
+            cls_map[msg.type](content=msg.content)
+            for msg in state["messages"][1:]
+        ]
+        translated = [SystemMessage(content=reflection_prompt)] + translated
+        res = self.llm.invoke(translated, {"configurable": {"thread_id": self.thread_id}})
+        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)
+
+        self.graph.add_edge(START, "generate")
+        self.graph.add_edge("generate", "reflect")
+        self.graph.add_edge("formalize", END)
+
+        self.graph.add_conditional_edges(
+            "reflect",
+            should_continue,
+            {"generate": "generate", "formalize": "formalize"},
+        )
+
+        # memory      = MemorySaver()
+        # self.action = self.graph.compile(checkpointer=memory)
+        self.action = self.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}})
+
+
+config = {"configurable": {"thread_id": "1"}}
+
+
+def should_continue(state: PlanningState):
+    if len(state["messages"]) > (state.get("reflection_steps", 3) + 3):
+        return "formalize"
+    if "[APPROVED]" in state["messages"][-1].content:
+        return "formalize"
+    return "generate"
+
+
+def main():
+    planning_agent = PlanningAgent()
+    for event in planning_agent.action.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."
+                )
+            ],
+        },
+        config,
+    ):
+        print("-" * 30)
+        print(event.keys())
+        print(event[list(event.keys())[0]]["messages"][-1].content)
+        print("-" * 30)
+
+
+if __name__ == "__main__":
+    main()

ursa/agents/recall_agent.py

@@ -0,0 +1,25 @@
+import os
+from .base import BaseAgent
+
+
+class RecallAgent(BaseAgent):
+    def __init__(self, llm, memory, **kwargs):
+        
+        super().__init__(llm, **kwargs)
+        self.memorydb = memory
+    
+    def remember(self, query):
+        memories = self.memorydb.retrieve(query)
+        summarize_query = f"""
+        You are being given the critical task of generating a detailed description of logged information 
+        to an important official to make a decision. Summarize the following memories that are related to 
+        the statement. Ensure that any specific details that are important are retained in the summary.
+
+        Query: {query}
+
+        """
+
+        for memory in memories:
+            summarize_query += f"Memory: {memory} \n\n"
+        memory = self.llm.invoke(summarize_query).content
+        return memory

ursa/agents/websearch_agent.py

@@ -0,0 +1,193 @@
+import inspect
+
+# from langchain_community.tools    import TavilySearchResults
+# from langchain_core.runnables.graph import MermaidDrawMethod
+from typing import Annotated, Any, List, Optional
+
+import requests
+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
+from langgraph.prebuilt import InjectedState, create_react_agent
+from pydantic import Field
+from typing_extensions import TypedDict
+
+from ..prompt_library.websearch_prompts import (
+    reflection_prompt,
+    websearch_prompt,
+    summarize_prompt,
+)
+from .base import BaseAgent
+
+# --- ANSI color codes ---
+BLUE = "\033[1;34m"
+RED = "\033[1;31m"
+GREEN = "\033[92m"
+RESET = "\033[0m"
+
+
+class WebSearchState(TypedDict):
+    websearch_query: str
+    messages: Annotated[list, add_messages]
+    urls_visited: List[str]
+    max_websearch_steps: Optional[int] = Field(
+        default=100, description="Maximum number of websearch steps"
+    )
+    remaining_steps: int
+    is_last_step: bool
+    model: Any
+
+
+# Adding the model to the state clumsily so that all "read" sources arent in the
+# context window. That eats a ton of tokens because each `llm.invoke` passes
+# all the tokens of all the sources.
+
+
+class WebSearchAgent(BaseAgent):
+    def __init__(
+        self, llm: str | BaseChatModel = "openai/gpt-4o-mini", **kwargs
+    ):
+        super().__init__(llm, **kwargs)
+        self.websearch_prompt = websearch_prompt
+        self.reflection_prompt = reflection_prompt
+        self.tools = [search_tool, process_content]  # + cb_tools
+        self.has_internet = self._check_for_internet(kwargs.get("url","http://www.lanl.gov"))
+        self._initialize_agent()
+
+    def review_node(self, state: WebSearchState) -> WebSearchState:
+        if not self.has_internet:
+            return {"messages":[HumanMessage(content="No internet for WebSearch Agent so no research to review.")], "urls_visited": []}
+        
+        translated = [SystemMessage(content=reflection_prompt)] + state[
+            "messages"
+        ]
+        res = self.llm.invoke(translated, {"configurable": {"thread_id": self.thread_id}})
+        return {"messages": [HumanMessage(content=res.content)]}
+
+    def response_node(self, state: WebSearchState) -> WebSearchState:
+        if not self.has_internet:
+            return {"messages":[HumanMessage(content="No internet for WebSearch Agent. No research carried out.")], "urls_visited": []}
+
+        messages = state["messages"] + [SystemMessage(content=summarize_prompt)]
+        response = self.llm.invoke(messages, {"configurable": {"thread_id": self.thread_id}})
+
+        urls_visited = []
+        for message in messages:
+            if message.model_dump().get("tool_calls", []):
+                if "url" in message.tool_calls[0]["args"]:
+                    urls_visited.append(message.tool_calls[0]["args"]["url"])
+        return {"messages": [response.content], "urls_visited": urls_visited}
+    
+    def _check_for_internet(self, url, timeout=2):
+        """
+        Checks for internet connectivity by attempting an HTTP GET request.
+        """
+        try:
+            requests.get(url, timeout=timeout)
+            return True
+        except (requests.ConnectionError, requests.Timeout):
+            return False
+
+    def _initialize_agent(self):
+        self.graph = StateGraph(WebSearchState)
+        self.graph.add_node(
+            "websearch",
+            create_react_agent(
+                self.llm,
+                self.tools,
+                state_schema=WebSearchState,
+                prompt=self.websearch_prompt,
+            ),
+        )
+
+        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("websearch", "review")
+        self.graph.add_edge("response", END)
+
+        self.graph.add_conditional_edges(
+            "review",
+            should_continue,
+            {"websearch": "websearch", "response": "response"},
+        )
+        self.action = self.graph.compile(checkpointer=self.checkpointer)
+        # self.action.get_graph().draw_mermaid_png(output_file_path="./websearch_agent_graph.png", draw_method=MermaidDrawMethod.PYPPETEER)
+    
+    def run(self, prompt, recursion_limit=100):
+        if not self.has_internet:
+            return {"messages":[HumanMessage(content="No internet for WebSearch Agent. No research carried out.")]}
+        inputs = {
+            "messages": [HumanMessage(content=prompt)],
+            "model": self.llm,
+        }
+        return self.action.invoke(inputs, {"recursion_limit":recursion_limit, "configurable": {"thread_id": self.thread_id}})        
+
+def process_content(
+    url: str, context: str, state: Annotated[dict, InjectedState]
+) -> str:
+    """
+    Processes content from a given webpage.
+
+    Args:
+        url: string with the url to obtain text content from.
+        context: string summary of the information the agent wants from the url for summarizing salient information.
+    """
+    print("Parsing information from ", url)
+    response = requests.get(url)
+    soup = BeautifulSoup(response.content, "html.parser")
+
+    content_prompt = f"""
+    Here is the full content:
+    {soup.get_text()}
+
+    Carefully summarize the content in full detail, given the following context:
+    {context}
+    """
+    summarized_information = state["model"].invoke(content_prompt, {"configurable": {"thread_id": self.thread_id}}).content
+    return summarized_information
+
+
+search_tool = DuckDuckGoSearchResults(output_format="json", num_results=10)
+# search_tool = TavilySearchResults(max_results=10, search_depth="advanced",include_answer=True)
+
+
+def should_continue(state: WebSearchState):
+    if len(state["messages"]) > (state.get("max_websearch_steps", 100) + 3):
+        return "response"
+    if "[APPROVED]" in state["messages"][-1].content:
+        return "response"
+    return "websearch"
+
+
+def main():
+    model = ChatOpenAI(
+        model="gpt-4o", max_tokens=10000, timeout=None, max_retries=2
+    )
+    websearcher = WebSearchAgent(llm=model)
+    problem_string = "Who are the 2025 Detroit Tigers top 10 prospects and what year were they born?"
+    inputs = {
+        "messages": [HumanMessage(content=problem_string)],
+        "model": model,
+    }
+    result = websearcher.action.invoke(inputs, {"recursion_limit": 10000, "configurable": {"thread_id": self.thread_id}})
+
+    colors = [BLUE, RED]
+    for ii, x in enumerate(result["messages"][:-1]):
+        if not isinstance(x, ToolMessage):
+            print(f"{colors[ii % 2]}" + x.content + f"{RESET}")
+
+    print(80 * "#")
+    print(f"{GREEN}" + result["messages"][-1].content + f"{RESET}")
+    print("Citations: ", result["urls_visited"])
+    return result
+
+
+if __name__ == "__main__":
+    main()

ursa/prompt_library/code_review_prompts.py

@@ -0,0 +1,51 @@
+
+def get_code_review_prompt(project_prompt, file_list):
+   return f'''
+   You are a responsible and efficient code review agent tasked with assessing if given files meet the goals of a project description.
+
+   The project goals are:
+   {project_prompt}
+
+   The list of files is:
+   {file_list}
+
+   Your responsibilities are as follows:
+
+   1. Read and review the given file and assess if it meets its requirements.
+      - Do not trust that the contents of the file are reflected in the filename without checking.
+   2. Ensure that all code uses real data and fully addresses the problem
+      - No fake, synthetic, or placeholder data. Obtain any needed data by reliable means.
+      - No simplifying assumptions. Ensure that code is implemented at the fully complexity required.
+      - Remove any code that may be dangerous, adversarial, or performing actions detrimental to the plan.
+      - Ensure files work together modularly, do not duplicate effort! 
+      - The project code should be clean and results reproducible.
+   3. Clearly document each action you take, including:
+      - The tools or methods you used.
+      - Any changes made including where the change occurred.
+      - Outcomes, results, or errors encountered during execution.
+   4. Immediately highlight and clearly communicate any steps that appear unclear, unsafe, or impractical before proceeding.
+
+   Your goal is to ensure the implemented code addresses the plan accurately, safely, and transparently, maintaining accountability at each step.
+   '''
+
+def get_plan_review_prompt(project_prompt, file_list):
+   return f'''
+   You are a responsible and efficient code review agent tasked with assessing if given files meet the goals of a project description.
+
+   The project goals are:
+   {project_prompt}
+
+   The list of files is:
+   {file_list}
+
+   Your responsibilities are as follows:
+
+   1. Formulate how the list of files work together to solve the given problem.
+   2. List potential problems to be reviewed in each file:
+      - Is any work duplicated or are the steps properly modularized?
+      - Does the file organization reflect a clear, reproducible workflow?
+      - Are there extraneous files or missing steps?
+      - Do any files appear dangerous, adversarial, or performing actions detrimental to the plan.
+
+   Your goal is to provide that information in a clear, concise way for use by a code reviewer who will look over files in detail.
+   '''