aiagents4pharma 1.39.4__py3-none-any.whl → 1.40.0__py3-none-any.whl

aiagents4pharma/talk2aiagents4pharma/configs/agents/main_agent/default.yaml

@@ -1,16 +1,29 @@
 _target_: agents.main_agent.get_app
 system_prompt: >
-  You are Talk2AIAgents4Pharma agent.
-  You are managing a team of the following 2 agents:
-   
-  1. Talk2Biomodels (T2B) agent: This agent can search and 
-  operate on mathematical models of biological systems. This
-  agent can also query an uploaded document/pdf/article. 
-   
-  2. Talk2KnowledgeGraphs (T2KG) agent: This agent can
-  reason over a knowledge graph of biological entities
-  and their relationships to answer subjective questions.
-  Do not call this agent for objective questions or questions
-  related to the mathematical models of biological systems.
+  You are the Talk2AIAgents4Pharma main agent. 
+  Your primary responsibility is to analyze the user's query and intelligently route it to the most appropriate specialized agent based on the user's intent.
 
-  Below mentioned are the prompts for each agent:
+  You manage the following two expert agents:
+
+  1. Talk2Biomodels (T2B) Agent:
+     - This agent specializes in working with **mathematical models of biological systems**.
+     - It can perform tasks such as querying and simulating models, conducting parameter scans, analyzing time series, 
+     and accessing annotated model components like species and reactions.
+     - It is also capable of processing and answering questions based on **uploaded documents, PDFs, or articles**.
+     - You should route queries to T2B if they contain or imply topics such as:
+       **"biological models", "mathematical models", "parameter scanning", "simulations", "time series", "species", "reactions", "annotations", "SBML", "model dynamics", "ODEs", "system biology"**, etc.
+     - For annotation-related queries, this agent should consider the annotated species and provide insights based on the retrieved context.
+
+  2. Talk2KnowledgeGraphs (T2KG) Agent:
+     - This agent specializes in **reasoning over biomedical knowledge graphs**.
+     - It is capable of answering **subjective, inferential, or relational questions** that involve **nodes, edges, relationships, and graph structures**.
+     - You should route queries to T2KG if they involve or imply topics such as:
+       **"knowledge graph", "biomedical graph", "subgraph", "graph structure", "node relationships", "graph reasoning", "graph connections", "pathways", "entity linkage", "graph queries", "subgraph extraction", "subgraph summarization"**, etc.
+
+  Important Guidelines:
+  - **Think carefully about the user's true intent**. Keywords are helpful indicators, but you must use reasoning to determine the best fit.
+  - If the question involves **mathematical modeling**, **simulation**, or **objective factual analysis**, select the **Talk2Biomodels (T2B)** agent.
+  - If the question involves **graph-based reasoning**, **subjective inference**, or **relationship analysis between entities**, select the **Talk2KnowledgeGraphs (T2KG)** agent.
+  - **Never** route a question about **biological models** or **objective computational simulations** to T2KG.
+
+  The following are the respective agent prompts:

aiagents4pharma/talk2aiagents4pharma/tests/test_main_agent.py

@@ -1,10 +1,11 @@
 '''
 Test Talk2AIAgents4Pharma supervisor agent.
 '''
-
+from unittest.mock import patch, MagicMock
 import pytest
 from langchain_core.messages import HumanMessage
 from langchain_openai import ChatOpenAI, OpenAIEmbeddings
+import pandas as pd
 from ..agents.main_agent import get_app
 
 # Define the data path for the test files of Talk2KnowledgeGraphs agent
@@ -41,6 +42,47 @@ def input_dict_fixture():
 
     return input_dict
 
+def mock_milvus_collection(name):
+    """
+    Mock Milvus collection for testing.
+    """
+    nodes = MagicMock()
+    nodes.query.return_value = [
+        {"node_index": 0,
+         "node_id": "id1",
+         "node_name": "Adalimumab",
+         "node_type": "drug",
+         "feat": "featA", "feat_emb": [0.1, 0.2, 0.3],
+         "desc": "descA", "desc_emb": [0.1, 0.2, 0.3]},
+        {"node_index": 1,
+         "node_id": "id2",
+         "node_name": "TNF",
+         "node_type": "gene/protein",
+         "feat": "featB", "feat_emb": [0.4, 0.5, 0.6], 
+         "desc": "descB", "desc_emb": [0.4, 0.5, 0.6]}
+    ]
+    nodes.load.return_value = None
+
+    edges = MagicMock()
+    edges.query.return_value = [
+        {"triplet_index": 0,
+         "head_id": "id1",
+         "head_index": 0,
+         "tail_id": "id2",
+         "tail_index": 1,
+         "edge_type": "drug,acts_on,gene/protein",
+         "display_relation": "acts_on",
+         "feat": "featC",
+         "feat_emb": [0.7, 0.8, 0.9]}
+    ]
+    edges.load.return_value = None
+
+    if "nodes" in name:
+        return nodes
+    if "edges" in name:
+        return edges
+    return None
+
 def test_main_agent_invokes_t2kg(input_dict):
     """
     In the following test, we will ask the main agent (supervisor)
@@ -67,8 +109,42 @@ def test_main_agent_invokes_t2kg(input_dict):
     )
     prompt = "List drugs that target the gene Interleukin-6"
 
-    # Invoke the agent
-    response = app.invoke({"messages": [HumanMessage(content=prompt)]}, config=config)
+    with patch("aiagents4pharma.talk2knowledgegraphs.tools."
+               "milvus_multimodal_subgraph_extraction.Collection", 
+               side_effect=mock_milvus_collection), \
+         patch("aiagents4pharma.talk2knowledgegraphs.tools."
+               "milvus_multimodal_subgraph_extraction.MultimodalPCSTPruning") as mock_pcst, \
+         patch("pymilvus.connections") as mock_connections, \
+         patch("aiagents4pharma.talk2knowledgegraphs.tools."
+               "milvus_multimodal_subgraph_extraction.hydra.initialize"), \
+         patch("aiagents4pharma.talk2knowledgegraphs.tools."
+               "milvus_multimodal_subgraph_extraction.hydra.compose") as mock_compose:
+        mock_connections.has_connection.return_value = True
+        mock_pcst_instance = MagicMock()
+        mock_pcst_instance.extract_subgraph.return_value = {
+            "nodes": pd.Series([0, 1]),
+            "edges": pd.Series([0])
+        }
+        mock_pcst.return_value = mock_pcst_instance
+        mock_cfg = MagicMock()
+        mock_cfg.cost_e = 1.0
+        mock_cfg.c_const = 1.0
+        mock_cfg.root = 0
+        mock_cfg.num_clusters = 1
+        mock_cfg.pruning = True
+        mock_cfg.verbosity_level = 0
+        mock_cfg.search_metric_type = "L2"
+        mock_cfg.node_colors_dict = {"drug": "blue", "gene/protein": "red"}
+        mock_compose.return_value = MagicMock()
+        mock_compose.return_value.tools.multimodal_subgraph_extraction = mock_cfg
+        mock_compose.return_value.tools.subgraph_summarization.\
+            prompt_subgraph_summarization = (
+            "Summarize the following subgraph: {textualized_subgraph}"
+        )
+
+        # Invoke the agent
+        response = app.invoke({"messages": [HumanMessage(content=prompt)]},
+                              config=config)
 
     # Check assistant message
     assistant_msg = response["messages"][-1].content
@@ -88,6 +164,10 @@ def test_main_agent_invokes_t2kg(input_dict):
     # Check summarized subgraph
     assert isinstance(dic_extracted_graph["graph_summary"], str)
 
+    # Another test for unknown collection
+    result = mock_milvus_collection("unknown")
+    assert result is None
+
 def test_main_agent_invokes_t2b():
     '''
     In the following test, we will ask the main agent (supervisor)

aiagents4pharma/talk2knowledgegraphs/agents/t2kg_agent.py

@@ -9,7 +9,10 @@ from langchain_core.language_models.chat_models import BaseChatModel
 from langgraph.checkpoint.memory import MemorySaver
 from langgraph.graph import START, StateGraph
 from langgraph.prebuilt import create_react_agent, ToolNode, InjectedState
-from ..tools.multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
+# from ..tools.multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
+from ..tools.milvus_multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
+# from ..tools.cu2_multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
+# from ..tools.gsfs_multimodal_subgraph_extraction import MultimodalSubgraphExtractionTool
 from ..tools.subgraph_summarization import SubgraphSummarizationTool
 from ..tools.graphrag_reasoning import GraphRAGReasoningTool
 from ..states.state_talk2knowledgegraphs import Talk2KnowledgeGraphs

aiagents4pharma/talk2knowledgegraphs/configs/app/frontend/default.yaml

@@ -6,7 +6,16 @@ multimodal_allowed_file_types:
   - "xls"
   - "xlsx"
 upload_data_dir: "../files"
-kg_name: "PrimeKG"
+kg_name: "BioBridge-PrimeKG"
+kg_node_types:
+  - "gene/protein"
+  - "molecular_function"
+  - "cellular_component"
+  - "biological_process"
+  - "drug"
+  - "disease"
+# kg_nodes_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/biobridge_nodes.parquet.gzip"
+# kg_edges_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/biobridge_edges.parquet.gzip"
 kg_pyg_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/biobridge_multimodal_pyg_graph.pkl"
 kg_text_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/biobridge_multimodal_text_graph.pkl"
 openai_api_key: ${oc.env:OPENAI_API_KEY}
@@ -15,6 +24,7 @@ openai_llms:
   - "gpt-4-turbo"
   - "gpt-3.5-turbo"
 openai_embeddings:
+  - "text-embedding-ada-002"
   - "text-embedding-3-small"
 ollama_llms:
   - "llama3.2:1b"
@@ -22,11 +32,32 @@ ollama_llms:
   - "llama3.1"
 ollama_embeddings:
   - "nomic-embed-text"
+default_embedding_model: "openai"
 temperature: 0.1
 streaming: False
-reasoning_subgraph_topk_nodes: 5
+reasoning_subgraph_topk_nodes: 15
 reasoning_subgraph_topk_nodes_min: 1
-reasoning_subgraph_topk_nodes_max: 100
-reasoning_subgraph_topk_edges: 5
+reasoning_subgraph_topk_nodes_max: 50
+reasoning_subgraph_topk_edges: 15
 reasoning_subgraph_topk_edges_min: 1
-reasoning_subgraph_topk_edges_max: 100
+reasoning_subgraph_topk_edges_max: 50
+# Configuration for Milvus
+milvus_db:
+  alias: "default"
+  host: ${oc.env:MILVUS_HOST,localhost}
+  port: ${oc.env:MILVUS_PORT,19530}
+  uri: "http://${oc.env:MILVUS_HOST,localhost}:${oc.env:MILVUS_PORT,19530}"
+  token: "root:Milvus"
+  user: "root"
+  password: "Milvus"
+  database_name: "t2kg_primekg"
+  collection_edges: "t2kg_primekg_edges"
+  collection_nodes: "t2kg_primekg_nodes"
+  collection_nodes_gene_protein: "t2kg_primekg_nodes_gene_protein"
+  collection_nodes_molecular_function: "t2kg_primekg_nodes_molecular_function"
+  collection_nodes_cellular_component: "t2kg_primekg_nodes_cellular_component"
+  collection_nodes_biological_process: "t2kg_primekg_nodes_biological_process"
+  collection_nodes_drug: "t2kg_primekg_nodes_drug"
+  collection_nodes_disease: "t2kg_primekg_nodes_disease"
+  query_batch_size: 1000000
+  cache_edge_index_path: "aiagents4pharma/talk2knowledgegraphs/tests/files/t2kg_primekg_edge_index.pkl"