aiagents4pharma 1.11.0__py3-none-any.whl → 1.12.0__py3-none-any.whl

aiagents4pharma/talk2knowledgegraphs/__init__.py

@@ -1,4 +1,5 @@
 '''
-This file is used to import the datasets, utils, and tools.
+This file is used to import the datasets and utils.
 '''
 from . import datasets
+from . import utils

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_enrichments.py

@@ -0,0 +1,39 @@
+"""
+Test cases for utils/enrichments/enrichments.py
+"""
+
+from ..utils.enrichments.enrichments import Enrichments
+
+class TestEnrichments(Enrichments):
+    """Test implementation of the Enrichments interface for testing purposes."""
+
+    def enrich_documents(self, texts: list[str]) -> list[list[float]]:
+        return [
+            f"Additional text description of {text} as the input." for text in texts
+        ]
+
+    def enrich_documents_with_rag(self, texts, docs):
+        # Currently we don't have a RAG model to test this method.
+        # Thus, we will just call the enrich_documents method instead.
+        return self.enrich_documents(texts)
+
+def test_enrich_documents():
+    """Test enriching documents using the Enrichments interface."""
+    enrichments = TestEnrichments()
+    texts = ["text1", "text2"]
+    result = enrichments.enrich_documents(texts)
+    assert result == [
+        "Additional text description of text1 as the input.",
+        "Additional text description of text2 as the input.",
+    ]
+
+def test_enrich_documents_with_rag():
+    """Test enriching documents with RAG using the Enrichments interface."""
+    enrichments = TestEnrichments()
+    texts = ["text1", "text2"]
+    docs = ["doc1", "doc2"]
+    result = enrichments.enrich_documents_with_rag(texts, docs)
+    assert result == [
+        "Additional text description of text1 as the input.",
+        "Additional text description of text2 as the input.",
+    ]

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py

@@ -0,0 +1,117 @@
+"""
+Test cases for utils/enrichments/ollama.py
+"""
+
+import pytest
+import ollama
+from ..utils.enrichments.ollama import EnrichmentWithOllama
+
+@pytest.fixture(name="ollama_config")
+def fixture_ollama_config():
+    """Return a dictionary with Ollama configuration."""
+    return {
+        "model_name": "smollm2:360m",
+        "prompt_enrichment": """
+            Given the input as a list of strings, please return the list of addditional information of
+            each input terms using your prior knowledge.
+
+            Example:
+            Input: ['acetaminophen', 'aspirin']
+            Ouput: ['acetaminophen is a medication used to treat pain and fever', 
+            'aspirin is a medication used to treat pain, fever, and inflammation']
+            
+            Do not include any pretext as the output, only the list of strings enriched.
+
+            Input: {input}
+        """,
+        "temperature": 0.0,
+        "streaming": False,
+    }
+
+def test_no_model_ollama(ollama_config):
+    """Test the case when the Ollama model is not available."""
+    cfg = ollama_config
+    cfg_model = "smollm2:135m" # Choose a small model
+
+    # Delete the Ollama model
+    try:
+        ollama.delete(cfg_model)
+    except ollama.ResponseError:
+        pass
+
+    # Check if the model is available
+    with pytest.raises(
+        ValueError, match=f"Error: Pulled {cfg_model} model and restarted Ollama server."
+    ):
+        EnrichmentWithOllama(
+            model_name=cfg_model,
+            prompt_enrichment=cfg["prompt_enrichment"],
+            temperature=cfg["temperature"],
+            streaming=cfg["streaming"],
+        )
+    ollama.delete(cfg_model)
+
+def test_enrich_nodes_ollama(ollama_config):
+    """Test the Ollama textual enrichment class for node enrichment."""
+    # Prepare enrichment model
+    cfg = ollama_config
+    enr_model = EnrichmentWithOllama(
+        model_name=cfg["model_name"],
+        prompt_enrichment=cfg["prompt_enrichment"],
+        temperature=cfg["temperature"],
+        streaming=cfg["streaming"],
+    )
+
+    # Perform enrichment for nodes
+    nodes = ["Adalimumab", "Infliximab"]
+    enriched_nodes = enr_model.enrich_documents(nodes)
+    # Check the enriched nodes
+    assert len(enriched_nodes) == 2
+    assert all(
+        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
+    )
+
+
+def test_enrich_relations_ollama(ollama_config):
+    """Test the Ollama textual enrichment class for relation enrichment."""
+    # Prepare enrichment model
+    cfg = ollama_config
+    enr_model = EnrichmentWithOllama(
+        model_name=cfg["model_name"],
+        prompt_enrichment=cfg["prompt_enrichment"],
+        temperature=cfg["temperature"],
+        streaming=cfg["streaming"],
+    )
+    # Perform enrichment for relations
+    relations = [
+        "IL23R-gene causation disease-inflammatory bowel diseases",
+        "NOD2-gene causation disease-inflammatory bowel diseases",
+    ]
+    enriched_relations = enr_model.enrich_documents(relations)
+    # Check the enriched relations
+    assert len(enriched_relations) == 2
+    assert all(
+        enriched_relations[i] != relations[i]
+        for i in range(len(relations))
+    )
+
+
+def test_enrich_ollama_rag(ollama_config):
+    """Test the Ollama textual enrichment class for enrichment with RAG (not implemented)."""
+    # Prepare enrichment model
+    cfg = ollama_config
+    enr_model = EnrichmentWithOllama(
+        model_name=cfg["model_name"],
+        prompt_enrichment=cfg["prompt_enrichment"],
+        temperature=cfg["temperature"],
+        streaming=cfg["streaming"],
+    )
+    # Perform enrichment for nodes
+    nodes = ["Adalimumab", "Infliximab"]
+    docs = [r"\path\to\doc1", r"\path\to\doc2"]
+    enriched_nodes = enr_model.enrich_documents_with_rag(nodes, docs)
+    # Check the enriched nodes
+    assert len(enriched_nodes) == 2
+    assert all(
+        enriched_nodes[i] != nodes[i] for i in range(len(nodes))
+    )

aiagents4pharma/talk2knowledgegraphs/utils/__init__.py

@@ -0,0 +1,5 @@
+'''
+This file is used to import utlities.
+'''
+from . import enrichments
+from . import embeddings

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py

@@ -0,0 +1,5 @@
+"""
+This package contains modules to use the enrichment model
+"""
+from . import enrichments
+from . import ollama

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/enrichments.py

@@ -0,0 +1,36 @@
+"""
+Enrichments interface
+"""
+
+from abc import ABC, abstractmethod
+
+class Enrichments(ABC):
+    """Interface for enrichment models.
+
+    This is an interface meant for implementing text enrichment models.
+
+    Enrichment models are used to enrich node or relation features in a given knowledge graph.
+    """
+
+    @abstractmethod
+    def enrich_documents(self, texts: list[str]) -> list[list[str]]:
+        """Enrich documents.
+
+        Args:
+            texts: List of documents to enrich.
+
+        Returns:
+            List of enriched documents.
+        """
+
+    @abstractmethod
+    def enrich_documents_with_rag(self, texts: list[str], docs: list[str]) -> list[str]:
+        """Enrich documents with RAG.
+
+        Args:
+            texts: List of documents to enrich.
+            docs: List of reference documents to enrich the input texts.
+
+        Returns:
+            List of enriched documents with RAG.
+        """

aiagents4pharma/talk2knowledgegraphs/utils/enrichments/ollama.py

@@ -0,0 +1,123 @@
+#!/usr/bin/env python3
+
+"""
+Enrichment class using Ollama model based on LangChain Enrichment class.
+"""
+
+import time
+from typing import List
+import subprocess
+import ast
+import ollama
+from langchain_ollama import ChatOllama
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from .enrichments import Enrichments
+
+class EnrichmentWithOllama(Enrichments):
+    """
+    Enrichment class using Ollama model based on the Enrichment abstract class.
+    """
+    def __init__(
+        self,
+        model_name: str,
+        prompt_enrichment: str,
+        temperature: float,
+        streaming: bool,
+    ):
+        """
+        Initialize the EnrichmentWithOllama class.
+
+        Args:
+            model_name: The name of the Ollama model to be used.
+            prompt_enrichment: The prompt enrichment template.
+            temperature: The temperature for the Ollama model.
+            streaming: The streaming flag for the Ollama model.
+        """
+        # Setup the Ollama server
+        self.__setup(model_name)
+
+        # Set parameters
+        self.model_name = model_name
+        self.prompt_enrichment = prompt_enrichment
+        self.temperature = temperature
+        self.streaming = streaming
+
+        # Prepare prompt template
+        self.prompt_template = ChatPromptTemplate.from_messages(
+            [
+                ("system", self.prompt_enrichment),
+                ("human", "{input}"),
+            ]
+        )
+
+        # Prepare model
+        self.model = ChatOllama(
+            model=self.model_name,
+            temperature=self.temperature,
+            streaming=self.streaming,
+        )
+
+    def __setup(self, model_name: str) -> None:
+        """
+        Check if the Ollama model is available and run the Ollama server if needed.
+
+        Args:
+            model_name: The name of the Ollama model to be used.
+        """
+        try:
+            models_list = ollama.list()["models"]
+            if model_name not in [m['model'].replace(":latest", "") for m in models_list]:
+                ollama.pull(model_name)
+                time.sleep(30)
+                raise ValueError(f"Pulled {model_name} model")
+        except Exception as e:
+            with subprocess.Popen(
+                "ollama serve", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            ):
+                time.sleep(10)
+            raise ValueError(f"Error: {e} and restarted Ollama server.") from e
+
+    def enrich_documents(self, texts: List[str]) -> List[str]:
+        """
+        Enrich a list of input texts with additional textual features using OLLAMA model.
+        Important: Make sure the input is a list of texts based on the defined prompt template
+        with 'input' as the variable name.
+
+        Args:
+            texts: The list of texts to be enriched.
+
+        Returns:
+            The list of enriched texts.
+        """
+
+        # Perform enrichment
+        chain = self.prompt_template | self.model | StrOutputParser()
+
+        # Generate the enriched node
+        # Important: Make sure the input is a list of texts based on the defined prompt template
+        # with 'input' as the variable name
+        enriched_texts = chain.invoke({"input": "[" + ", ".join(texts) + "]"})
+
+        # Convert the enriched nodes to a list of dictionary
+        enriched_texts = ast.literal_eval(enriched_texts.replace("```", ""))
+
+        # Final check for the enriched texts
+        assert len(enriched_texts) == len(texts)
+
+        return enriched_texts
+
+    def enrich_documents_with_rag(self, texts, docs):
+        """
+        Enrich a list of input texts with additional textual features using OLLAMA model with RAG.
+        As of now, we don't have a RAG model to test this method yet.
+        Thus, we will just call the enrich_documents method instead.
+
+        Args:
+            texts: The list of texts to be enriched.
+            docs: The list of reference documents to enrich the input texts.
+        
+        Returns:
+            The list of enriched texts
+        """
+        return self.enrich_documents(texts)

{aiagents4pharma-1.11.0.dist-info → aiagents4pharma-1.12.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: aiagents4pharma
-Version: 1.11.0
+Version: 1.12.0
 Summary: AI Agents for drug discovery, drug development, and other pharmaceutical R&D
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -20,9 +20,11 @@ Requires-Dist: langchain-community==0.3.5
 Requires-Dist: langchain-core==0.3.31
 Requires-Dist: langchain-experimental==0.3.3
 Requires-Dist: langchain-openai==0.2.5
+Requires-Dist: langchain_ollama==0.2.2
 Requires-Dist: langgraph==0.2.66
 Requires-Dist: matplotlib==3.9.2
 Requires-Dist: openai==1.59.4
+Requires-Dist: ollama==0.4.6
 Requires-Dist: pandas==2.2.3
 Requires-Dist: plotly==5.24.1
 Requires-Dist: pydantic==2.9.2

{aiagents4pharma-1.11.0.dist-info → aiagents4pharma-1.12.0.dist-info}/RECORD

@@ -51,7 +51,7 @@ aiagents4pharma/talk2competitors/tools/s2/display_results.py,sha256=B8JJGohi1Eyx
 aiagents4pharma/talk2competitors/tools/s2/multi_paper_rec.py,sha256=FYLt47DAk6WOKfEk1Gj9zVvJGNyxA283PCp8IKW9U5M,4262
 aiagents4pharma/talk2competitors/tools/s2/search.py,sha256=pppjrQv5-8ep4fnqgTSBNgnbSnQsVIcNrRrH0p2TP1o,4025
 aiagents4pharma/talk2competitors/tools/s2/single_paper_rec.py,sha256=dAfUQxI7T5eu0eDxK8VAl7-JH0Wnw24CVkOQqwj-hXc,4810
-aiagents4pharma/talk2knowledgegraphs/__init__.py,sha256=SW7Ys2A4eXyFtizNPdSw91SHOPVUBGBsrCQ7TqwSUL0,91
+aiagents4pharma/talk2knowledgegraphs/__init__.py,sha256=4smVQoSMM6rflVnNkABqlDAAlSn4bYsq7rMVWjRGvis,103
 aiagents4pharma/talk2knowledgegraphs/datasets/__init__.py,sha256=L3gPuHskSegmtXskVrLIYr7FXe_ibKgJ2GGr1_Wok6k,173
 aiagents4pharma/talk2knowledgegraphs/datasets/biobridge_primekg.py,sha256=QlzDXmXREoa9MA6-GwzqRjdzndQeGBAF11Td6NFk_9Y,23426
 aiagents4pharma/talk2knowledgegraphs/datasets/dataset.py,sha256=-LaPLse8BkALqwFetNK7wch2dt9Dz6QKGKZKBKM6bIk,409
@@ -65,14 +65,19 @@ aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_starkqa_primekg.py,sha2
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_embeddings.py,sha256=uYFoE_6zeU10_1mLLAHUr5c4S2XZMSc0Q_860o-KWEw,1517
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_huggingface.py,sha256=EINWyXg_3AMHF3WzFLhIUiFDuaEhTVHBvVAJr8VtMDg,1624
 aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_sentencetransformer.py,sha256=Qxo6WeIDRy8aLh1tNKw0kSlzmUj3MtTak63oW2YwB24,1327
-aiagents4pharma/talk2knowledgegraphs/utils/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_enrichments.py,sha256=N6HRr4lWHXY7bTHe2uXJe4D_EG9WqZPibZne6qLl9_k,1447
+aiagents4pharma/talk2knowledgegraphs/tests/test_utils_enrichments_ollama.py,sha256=kMuB_vci6hKr2qJgXBmcje7yxeJ2nY2ImXw-NSJpts0,3912
+aiagents4pharma/talk2knowledgegraphs/utils/__init__.py,sha256=X8kSpmDTMkeE0fA5D0CWMsQ52YKoM5rRSXrjnali3IM,97
 aiagents4pharma/talk2knowledgegraphs/utils/kg_utils.py,sha256=6vQnPkeOWae_8jePjhma3sJuMTngy0I0tqzdFt6OqKg,2507
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/__init__.py,sha256=xRb0x7SoAb0nSVZYgjrqxWvENOMDuqIdL43NMjoOaCs,153
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/embeddings.py,sha256=1nGznrAj-xT0xuSMBGz2dOujJ7M_IwSR84njxtxsy9A,2523
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/huggingface.py,sha256=2vi_elf6EgzfagFAO5QnL3a_aXZyN7B1EBziu44MTfM,3806
 aiagents4pharma/talk2knowledgegraphs/utils/embeddings/sentence_transformer.py,sha256=36iKlisOpMtGR5xfTAlSHXWvPqVC_Jbezod8kbBBMVg,2136
-aiagents4pharma-1.11.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
-aiagents4pharma-1.11.0.dist-info/METADATA,sha256=GsvCuC24bJ_wwQP8aevOB4Eai8WcUIcr9kmN3EwYP_Y,8541
-aiagents4pharma-1.11.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
-aiagents4pharma-1.11.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
-aiagents4pharma-1.11.0.dist-info/RECORD,,
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/__init__.py,sha256=tW426knki2DBIHcWyF_K04iMMdbpIn_e_TpPmTgz2dI,113
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/enrichments.py,sha256=Bx8x6zzk5614ApWB90N_iv4_Y_Uq0-KwUeBwYSdQMU4,924
+aiagents4pharma/talk2knowledgegraphs/utils/enrichments/ollama.py,sha256=8eoxR-VHo0G7ReQIwje7xEhE-SJlHdef7_wJRpnvFIc,4116
+aiagents4pharma-1.12.0.dist-info/LICENSE,sha256=IcIbyB1Hyk5ZDah03VNQvJkbNk2hkBCDqQ8qtnCvB4Q,1077
+aiagents4pharma-1.12.0.dist-info/METADATA,sha256=Okx1RPc3qMmJ2a7ca2YvSFr7-gxqAR51LVLkOTh5O2k,8609
+aiagents4pharma-1.12.0.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+aiagents4pharma-1.12.0.dist-info/top_level.txt,sha256=-AH8rMmrSnJtq7HaAObS78UU-cTCwvX660dSxeM7a0A,16
+aiagents4pharma-1.12.0.dist-info/RECORD,,