ursa-ai 0.9.1__py3-none-any.whl

ursa/agents/rag_agent.py

@@ -0,0 +1,304 @@
+import os
+import re
+import statistics
+from functools import cached_property
+from threading import Lock
+from typing import Any, Mapping, TypedDict
+
+from langchain.chat_models import BaseChatModel
+from langchain.embeddings import Embeddings
+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 langchain_text_splitters import RecursiveCharacterTextSplitter
+from langgraph.graph import StateGraph
+from tqdm import tqdm
+
+from ursa.agents.base import BaseAgent
+
+
+class RAGMetadata(TypedDict):
+    k: int
+    num_results: int
+    relevance_scores: list[float]
+
+
+class RAGState(TypedDict, total=False):
+    context: str
+    doc_texts: list[str]
+    doc_ids: list[str]
+    summary: str
+    rag_metadata: RAGMetadata
+
+
+def remove_surrogates(text: str) -> str:
+    return re.sub(r"[\ud800-\udfff]", "", text)
+
+
+class RAGAgent(BaseAgent):
+    def __init__(
+        self,
+        embedding: Embeddings,
+        llm: BaseChatModel,
+        return_k: int = 10,
+        chunk_size: int = 1000,
+        chunk_overlap: int = 200,
+        database_path: str = "database",
+        summaries_path: str = "database",
+        vectorstore_path: str = "vectorstore",
+        **kwargs,
+    ):
+        super().__init__(llm, **kwargs)
+        self.retriever = None
+        self._vs_lock = Lock()
+        self.return_k = return_k
+        self.embedding = embedding
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.database_path = database_path
+        self.summaries_path = summaries_path
+        self.vectorstore_path = vectorstore_path
+
+        os.makedirs(self.vectorstore_path, exist_ok=True)
+        self.vectorstore = self._open_global_vectorstore()
+
+    @cached_property
+    def graph(self):
+        return self._build_graph()
+
+    @property
+    def _action(self):
+        return self.graph
+
+    @property
+    def manifest_path(self) -> str:
+        return os.path.join(self.vectorstore_path, "_ingested_ids.txt")
+
+    @property
+    def manifest_exists(self) -> bool:
+        return os.path.exists(self.manifest_path)
+
+    def _open_global_vectorstore(self) -> Chroma:
+        return Chroma(
+            persist_directory=self.vectorstore_path,
+            embedding_function=self.embedding,
+            collection_metadata={"hnsw:space": "cosine"},
+        )
+
+    def _paper_exists_in_vectorstore(self, doc_id: str) -> bool:
+        try:
+            col = self.vectorstore._collection
+            res = col.get(where={"id": doc_id}, limit=1)
+            return len(res.get("ids", [])) > 0
+        except Exception:
+            if not self.manifest_exists:
+                return False
+            with open(self.manifest_path, "r") as f:
+                return any(line.strip() == doc_id for line in f)
+
+    def _mark_paper_ingested(self, arxiv_id: str) -> None:
+        with open(self.manifest_path, "a") as f:
+            f.write(f"{arxiv_id}\n")
+
+    def _ensure_doc_in_vectorstore(self, paper_text: str, doc_id: str) -> None:
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
+        )
+        docs = splitter.create_documents(
+            [paper_text], metadatas=[{"id": doc_id}]
+        )
+        with self._vs_lock:
+            if not self._paper_exists_in_vectorstore(doc_id):
+                ids = [f"{doc_id}::{i}" for i, _ in enumerate(docs)]
+                self.vectorstore.add_documents(docs, ids=ids)
+                self._mark_paper_ingested(doc_id)
+
+    def _get_global_retriever(self, k: int = 5):
+        return self.vectorstore, self.vectorstore.as_retriever(
+            search_kwargs={"k": k}
+        )
+
+    def _read_docs_node(self, state: RAGState) -> RAGState:
+        print("[RAG Agent] Reading Documents....")
+        papers = []
+        new_state = state.copy()
+
+        pdf_files = [
+            f
+            for f in os.listdir(self.database_path)
+            if f.lower().endswith(".pdf")
+        ]
+
+        doc_ids = [
+            pdf_filename.rsplit(".pdf", 1)[0] for pdf_filename in pdf_files
+        ]
+        pdf_files = [
+            pdf_filename
+            for pdf_filename, id in zip(pdf_files, doc_ids)
+            if not self._paper_exists_in_vectorstore(id)
+        ]
+
+        for pdf_filename in tqdm(pdf_files, desc="RAG parsing text"):
+            full_text = ""
+
+            try:
+                loader = PyPDFLoader(
+                    os.path.join(self.database_path, pdf_filename)
+                )
+                pages = loader.load()
+                full_text = "\n".join([p.page_content for p in pages])
+
+            except Exception as e:
+                full_text = f"Error loading paper: {e}"
+
+            papers.append(full_text)
+
+        new_state["doc_texts"] = papers
+        new_state["doc_ids"] = doc_ids
+
+        return new_state
+
+    def _ingest_docs_node(self, state: RAGState) -> RAGState:
+        splitter = RecursiveCharacterTextSplitter(
+            chunk_size=self.chunk_size, chunk_overlap=self.chunk_overlap
+        )
+
+        if "doc_texts" not in state:
+            raise RuntimeError("Unexpected error: doc_ids not in state!")
+
+        if "doc_ids" not in state:
+            raise RuntimeError("Unexpected error: doc_texts not in state!")
+
+        batch_docs, batch_ids = [], []
+        for paper, id in tqdm(
+            zip(state["doc_texts"], state["doc_ids"]),
+            total=len(state["doc_texts"]),
+            desc="RAG Ingesting",
+        ):
+            cleaned_text = remove_surrogates(paper)
+            docs = splitter.create_documents(
+                [cleaned_text], metadatas=[{"id": id}]
+            )
+            ids = [f"{id}::{i}" for i, _ in enumerate(docs)]
+            batch_docs.extend(docs)
+            batch_ids.extend(ids)
+
+        if state["doc_texts"]:
+            print("[RAG Agent] Ingesting Documents Into RAG Database....")
+            with self._vs_lock:
+                self.vectorstore.add_documents(batch_docs, ids=batch_ids)
+                for id in batch_ids:
+                    self._mark_paper_ingested(id)
+
+        return state
+
+    def _retrieve_and_summarize_node(self, state: RAGState) -> RAGState:
+        print(
+            "[RAG Agent] Retrieving Contextually Relevant Information From Database..."
+        )
+        prompt = ChatPromptTemplate.from_template("""
+        You are a scientific assistant responsible for summarizing extracts from research papers, in the context of the following task: {context}
+
+        Summarize the retrieved scientific content below.
+        Cite sources by ID when relevant: {source_ids}
+                                                  
+        {retrieved_content}
+        """)
+        chain = prompt | self.llm | StrOutputParser()
+
+        # 2) One retrieval over the global DB with the task context
+        try:
+            if "context" not in state:
+                raise RuntimeError("Unexpected error: context not in state!")
+
+            results = self.vectorstore.similarity_search_with_relevance_scores(
+                state["context"], k=self.return_k
+            )
+
+            relevance_scores = [score for _, score in results]
+        except Exception as e:
+            print(f"RAG failed due to: {e}")
+            return {**state, "summary": ""}
+
+        source_ids_list = []
+        for doc, _ in results:
+            aid = doc.metadata.get("id")
+            if aid and aid not in source_ids_list:
+                source_ids_list.append(aid)
+        source_ids = ", ".join(source_ids_list)
+
+        retrieved_content = (
+            "\n\n".join(doc.page_content for doc, _ in results)
+            if results
+            else ""
+        )
+
+        print("[RAG Agent] Summarizing Retrieved Information From Database...")
+        # 3) One summary based on retrieved chunks
+        rag_summary = chain.invoke({
+            "retrieved_content": retrieved_content,
+            "context": state["context"],
+            "source_ids": source_ids,
+        })
+
+        # Persist a single file for the batch (optional)
+        batch_name = "RAG_summary.txt"
+        os.makedirs(self.summaries_path, exist_ok=True)
+        with open(os.path.join(self.summaries_path, batch_name), "w") as f:
+            f.write(rag_summary)
+
+        # Diagnostics
+        if relevance_scores:
+            print(f"\nMax Relevance Score: {max(relevance_scores):.4f}")
+            print(f"Min Relevance Score: {min(relevance_scores):.4f}")
+            print(
+                f"Median Relevance Score: {statistics.median(relevance_scores):.4f}\n"
+            )
+        else:
+            print("\nNo RAG results retrieved (score list empty).\n")
+
+        # Return a single-element list by default (preferred)
+        return {
+            **state,
+            "summary": rag_summary,
+            "rag_metadata": {
+                "k": self.return_k,
+                "num_results": len(results),
+                "relevance_scores": relevance_scores,
+            },
+        }
+
+    def _invoke(
+        self, inputs: Mapping[str, Any], recursion_limit: int = 100000, **_
+    ):
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
+        return self._action.invoke(inputs, config)
+
+    def _build_graph(self):
+        graph = StateGraph(RAGState)
+
+        self.add_node(graph, self._read_docs_node)
+        self.add_node(graph, self._ingest_docs_node)
+        self.add_node(graph, self._retrieve_and_summarize_node)
+
+        graph.add_edge("_read_docs_node", "_ingest_docs_node")
+        graph.add_edge("_ingest_docs_node", "_retrieve_and_summarize_node")
+
+        graph.set_entry_point("_read_docs_node")
+        graph.set_finish_point("_retrieve_and_summarize_node")
+
+        return graph.compile(checkpointer=self.checkpointer)
+
+
+# NOTE: Run test in `tests/agents/test_rag_agent/test_rag_agent.py` via:
+#
+# pytest -s tests/agents/test_rag_agent
+#
+# OR
+#
+# uv run pytest -s tests/agents/test_rag_agent
+#
+# NOTE: You may need to `rm -rf workspace/rag-agent` to remove the vectorstore.

ursa/agents/recall_agent.py

@@ -0,0 +1,54 @@
+from typing import Any, Mapping, TypedDict
+
+from langchain.chat_models import BaseChatModel
+from langgraph.graph import StateGraph
+
+from .base import BaseAgent
+
+
+class RecallState(TypedDict):
+    query: str
+    memory: str
+
+
+class RecallAgent(BaseAgent):
+    def __init__(self, llm: BaseChatModel, memory, **kwargs):
+        super().__init__(llm, **kwargs)
+        self.memorydb = memory
+        self._action = self._build_graph()
+
+    def _remember(self, state: RecallState) -> str:
+        memories = self.memorydb.retrieve(state["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: {state["query"]}
+
+        """
+
+        for memory in memories:
+            summarize_query += f"Memory: {memory} \n\n"
+        state["memory"] = self.llm.invoke(summarize_query).content
+        return state
+
+    def _build_graph(self):
+        graph = StateGraph(RecallState)
+
+        self.add_node(graph, self._remember)
+        graph.set_entry_point("_remember")
+        graph.set_finish_point("_remember")
+        return graph.compile(checkpointer=self.checkpointer)
+
+    def _invoke(
+        self, inputs: Mapping[str, Any], recursion_limit: int = 100000, **_
+    ):
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
+        if "query" not in inputs:
+            raise ("'query' is a required argument")
+
+        output = self._action.invoke(inputs, config)
+        return output["memory"]

ursa/agents/websearch_agent.py

@@ -0,0 +1,196 @@
+# from langchain_community.tools    import TavilySearchResults
+# from langchain_core.runnables.graph import MermaidDrawMethod
+from typing import Annotated, Any, Mapping
+
+import requests
+from bs4 import BeautifulSoup
+from langchain.agents import create_agent
+from langchain.chat_models import BaseChatModel
+from langchain.messages import HumanMessage, SystemMessage
+from langchain_community.tools import DuckDuckGoSearchResults
+from langgraph.graph import StateGraph
+from langgraph.graph.message import add_messages
+from langgraph.prebuilt import InjectedState
+from pydantic import Field
+from typing_extensions import TypedDict
+
+from ..prompt_library.websearch_prompts import (
+    reflection_prompt,
+    summarize_prompt,
+    websearch_prompt,
+)
+from .base import BaseAgent
+
+
+class WebSearchState(TypedDict):
+    websearch_query: str
+    messages: Annotated[list, add_messages]
+    urls_visited: list[str]
+    max_websearch_steps: Annotated[
+        int, Field(default=100, description="Maximum number of websearch steps")
+    ]
+    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
+# context window. That eats a ton of tokens because each `llm.invoke` passes
+# all the tokens of all the sources.
+
+
+class WebSearchAgentLegacy(BaseAgent):
+    def __init__(
+        self,
+        llm: BaseChatModel,
+        **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._build_graph()
+
+    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 _state_store_node(self, state: WebSearchState) -> WebSearchState:
+        state["thread_id"] = self.thread_id
+        return state
+        # return dict(**state, thread_id=self.thread_id)
+
+    def _create_react(self, state: WebSearchState) -> WebSearchState:
+        react_agent = create_agent(
+            self.llm,
+            self.tools,
+            state_schema=WebSearchState,
+            system_prompt=self.websearch_prompt,
+        )
+        return react_agent.invoke(state)
+
+    def _build_graph(self):
+        graph = StateGraph(WebSearchState)
+        self.add_node(graph, self._state_store_node)
+        self.add_node(graph, self._create_react)
+        self.add_node(graph, self._review_node)
+        self.add_node(graph, self._response_node)
+
+        graph.set_entry_point("_state_store_node")
+        graph.add_edge("_state_store_node", "_create_react")
+        graph.add_edge("_create_react", "_review_node")
+        graph.set_finish_point("_response_node")
+
+        graph.add_conditional_edges(
+            "_review_node",
+            should_continue,
+            {
+                "_create_react": "_create_react",
+                "_response_node": "_response_node",
+            },
+        )
+        self._action = graph.compile(checkpointer=self.checkpointer)
+        # self._action.get_graph().draw_mermaid_png(output_file_path="./websearch_agent_graph.png", draw_method=MermaidDrawMethod.PYPPETEER)
+
+    def _invoke(
+        self, inputs: Mapping[str, Any], recursion_limit: int = 1000, **_
+    ):
+        config = self.build_config(
+            recursion_limit=recursion_limit, tags=["graph"]
+        )
+        return self._action.invoke(inputs, config)
+
+
+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": state["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_node"
+    if "[APPROVED]" in state["messages"][-1].content:
+        return "_response_node"
+    return "_create_react"