aiagents4pharma 1.8.0__py3-none-any.whl → 1.15.0__py3-none-any.whl

aiagents4pharma/talk2competitors/tools/s2/single_paper_rec.py

@@ -0,0 +1,141 @@
+#!/usr/bin/env python3
+
+"""
+This tool is used to return recommendations for a single paper.
+"""
+
+import logging
+from typing import Annotated, Any, Dict, Optional
+
+import pandas as pd
+import requests
+from langchain_core.messages import ToolMessage
+from langchain_core.tools import tool
+from langchain_core.tools.base import InjectedToolCallId
+from langgraph.types import Command
+from pydantic import BaseModel, Field
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+
+class SinglePaperRecInput(BaseModel):
+    """Input schema for single paper recommendation tool."""
+
+    paper_id: str = Field(
+        description="Semantic Scholar Paper ID to get recommendations for (40-character string)"
+    )
+    limit: int = Field(
+        default=2,
+        description="Maximum number of recommendations to return",
+        ge=1,
+        le=500,
+    )
+    year: Optional[str] = Field(
+        default=None,
+        description="Year range in format: YYYY for specific year, "
+        "YYYY- for papers after year, -YYYY for papers before year, or YYYY:YYYY for range",
+    )
+    tool_call_id: Annotated[str, InjectedToolCallId]
+    model_config = {"arbitrary_types_allowed": True}
+
+
+@tool(args_schema=SinglePaperRecInput)
+def get_single_paper_recommendations(
+    paper_id: str,
+    tool_call_id: Annotated[str, InjectedToolCallId],
+    limit: int = 2,
+    year: Optional[str] = None,
+) -> Dict[str, Any]:
+    """
+    Get paper recommendations based on a single paper.
+
+    Args:
+        paper_id (str): The Semantic Scholar Paper ID to get recommendations for.
+        tool_call_id (Annotated[str, InjectedToolCallId]): The tool call ID.
+        limit (int, optional): The maximum number of recommendations to return. Defaults to 2.
+        year (str, optional): Year range for papers.
+        Supports formats like "2024-", "-2024", "2024:2025". Defaults to None.
+
+    Returns:
+        Dict[str, Any]: The recommendations and related information.
+    """
+    logger.info("Starting single paper recommendations search.")
+
+    endpoint = (
+        f"https://api.semanticscholar.org/recommendations/v1/papers/forpaper/{paper_id}"
+    )
+    params = {
+        "limit": min(limit, 500),  # Max 500 per API docs
+        "fields": "paperId,title,abstract,year,authors,citationCount,url",
+        "from": "all-cs",  # Using all-cs pool as specified in docs
+    }
+
+    # Add year parameter if provided
+    if year:
+        params["year"] = year
+
+    response = requests.get(endpoint, params=params, timeout=10)
+    data = response.json()
+    papers = data.get("data", [])
+    response = requests.get(endpoint, params=params, timeout=10)
+    # print(f"API Response Status: {response.status_code}")
+    logging.info(
+        "API Response Status for recommendations of paper %s: %s",
+        paper_id,
+        response.status_code,
+    )
+    # print(f"Request params: {params}")
+    logging.info("Request params: %s", params)
+
+    data = response.json()
+    recommendations = data.get("recommendedPapers", [])
+
+    # Extract paper ID and title from recommendations
+    filtered_papers = {
+        paper["paperId"]: {
+            "Title": paper.get("title", "N/A"),
+            "Abstract": paper.get("abstract", "N/A"),
+            "Year": paper.get("year", "N/A"),
+            "Citation Count": paper.get("citationCount", "N/A"),
+            "URL": paper.get("url", "N/A"),
+            # "Publication Type": paper.get("publicationTypes", ["N/A"])[0]
+            # if paper.get("publicationTypes")
+            # else "N/A",
+            # "Open Access PDF": paper.get("openAccessPdf", {}).get("url", "N/A")
+            # if paper.get("openAccessPdf") is not None
+            # else "N/A",
+        }
+        for paper in recommendations
+        if paper.get("title") and paper.get("authors")
+    }
+
+    # Create a DataFrame for pretty printing
+    df = pd.DataFrame(filtered_papers)
+
+    # Format papers for state update
+    papers = [
+        f"Paper ID: {paper_id}\n"
+        f"Title: {paper_data['Title']}\n"
+        f"Abstract: {paper_data['Abstract']}\n"
+        f"Year: {paper_data['Year']}\n"
+        f"Citations: {paper_data['Citation Count']}\n"
+        f"URL: {paper_data['URL']}\n"
+        # f"Publication Type: {paper_data['Publication Type']}\n"
+        # f"Open Access PDF: {paper_data['Open Access PDF']}"
+        for paper_id, paper_data in filtered_papers.items()
+    ]
+
+    # Convert DataFrame to markdown table
+    markdown_table = df.to_markdown(tablefmt="grid")
+    logging.info("Search results: %s", papers)
+
+    return Command(
+        update={
+            "papers": filtered_papers,  # Now sending the dictionary directly
+            "messages": [
+                ToolMessage(content=markdown_table, tool_call_id=tool_call_id)
+            ],
+        }
+    )

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_datasets_biobridge_primekg.py

@@ -0,0 +1,242 @@
+"""
+Test cases for datasets/primekg_loader.py
+"""
+
+import os
+import shutil
+import pytest
+from ..datasets.biobridge_primekg import BioBridgePrimeKG
+
+# Remove the data folder for testing if it exists
+PRIMEKG_LOCAL_DIR = "../data/primekg_test/"
+LOCAL_DIR = "../data/biobridge_primekg_test/"
+shutil.rmtree(LOCAL_DIR, ignore_errors=True)
+
+@pytest.fixture(name="biobridge_primekg")
+def biobridge_primekg_fixture():
+    """
+    Fixture for creating an instance of PrimeKG.
+    """
+    return BioBridgePrimeKG(primekg_dir=PRIMEKG_LOCAL_DIR,
+                            local_dir=LOCAL_DIR)
+
+def test_download_primekg(biobridge_primekg):
+    """
+    Test the loading method of the BioBridge-PrimeKG class by downloading data from repository.
+    """
+    # Load BioBridge-PrimeKG data
+    biobridge_primekg.load_data()
+    primekg_nodes = biobridge_primekg.get_primekg().get_nodes()
+    primekg_edges = biobridge_primekg.get_primekg().get_edges()
+    biobridge_data_config = biobridge_primekg.get_data_config()
+    biobridge_emb_dict = biobridge_primekg.get_node_embeddings()
+    biobridge_triplets = biobridge_primekg.get_primekg_triplets()
+    biobridge_splits = biobridge_primekg.get_train_test_split()
+    biobridge_node_info = biobridge_primekg.get_node_info_dict()
+
+    # Check if the local directories exists
+    assert os.path.exists(biobridge_primekg.primekg_dir)
+    assert os.path.exists(biobridge_primekg.local_dir)
+    # Check if downloaded and processed files exist
+    # PrimeKG files
+    files = ["nodes.tab", "primekg_nodes.tsv.gz",
+             "edges.csv", "primekg_edges.tsv.gz"]
+    for file in files:
+        path = f"{biobridge_primekg.primekg_dir}/{file}"
+        assert os.path.exists(path)
+    # BioBridge data config
+    assert os.path.exists(f"{biobridge_primekg.local_dir}/data_config.json")
+    # BioBridge embeddings
+    files = [
+        "protein.pkl",
+        "mf.pkl",
+        "cc.pkl",
+        "bp.pkl",
+        "drug.pkl",
+        "disease.pkl",
+        "embedding_dict.pkl"
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/embeddings/{file}"
+        assert os.path.exists(path)
+    # BioBridge processed files
+    files = [
+        "protein.csv",
+        "mf.csv",
+        "cc.csv",
+        "bp.csv",
+        "drug.csv",
+        "disease.csv",
+        "triplet_full.tsv.gz",
+        "triplet_full_altered.tsv.gz",
+        "node_train.tsv.gz",
+        "triplet_train.tsv.gz",
+        "node_test.tsv.gz",
+        "triplet_test.tsv.gz",
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/processed/{file}"
+        assert os.path.exists(path)
+    # Check processed PrimeKG dataframes
+    # Nodes
+    assert primekg_nodes is not None
+    assert len(primekg_nodes) > 0
+    assert primekg_nodes.shape[0] == 129375
+    # Edges
+    assert primekg_edges is not None
+    assert len(primekg_edges) > 0
+    assert primekg_edges.shape[0] == 8100498
+    # Check processed BioBridge data config
+    assert biobridge_data_config is not None
+    assert len(biobridge_data_config) > 0
+    assert len(biobridge_data_config['node_type']) == 10
+    assert len(biobridge_data_config['relation_type']) == 18
+    assert len(biobridge_data_config['emb_dim']) == 6
+    # Check processed BioBridge embeddings
+    assert biobridge_emb_dict is not None
+    assert len(biobridge_emb_dict) > 0
+    assert len(biobridge_emb_dict) == 85466
+    # Check processed BioBridge triplets
+    assert biobridge_triplets is not None
+    assert len(biobridge_triplets) > 0
+    assert biobridge_triplets.shape[0] == 3904610
+    assert list(biobridge_splits.keys()) == ['train', 'node_train', 'test', 'node_test']
+    assert len(biobridge_splits['train']) == 3510930
+    assert len(biobridge_splits['node_train']) == 76486
+    assert len(biobridge_splits['test']) == 393680
+    assert len(biobridge_splits['node_test']) == 8495
+    # Check node info dictionary
+    assert list(biobridge_node_info.keys()) == ['gene/protein',
+                                                'molecular_function',
+                                                'cellular_component',
+                                                'biological_process',
+                                                'drug',
+                                                'disease']
+    assert len(biobridge_node_info['gene/protein']) == 19162
+    assert len(biobridge_node_info['molecular_function']) == 10966
+    assert len(biobridge_node_info['cellular_component']) == 4013
+    assert len(biobridge_node_info['biological_process']) == 27478
+    assert len(biobridge_node_info['drug']) == 6948
+    assert len(biobridge_node_info['disease']) == 44133
+
+
+def test_load_existing_primekg(biobridge_primekg):
+    """
+    Test the loading method of the BioBridge-PrimeKG class by loading existing data in local.
+    """
+    # Load BioBridge-PrimeKG data
+    biobridge_primekg.load_data()
+    primekg_nodes = biobridge_primekg.get_primekg().get_nodes()
+    primekg_edges = biobridge_primekg.get_primekg().get_edges()
+    biobridge_data_config = biobridge_primekg.get_data_config()
+    biobridge_emb_dict = biobridge_primekg.get_node_embeddings()
+    biobridge_triplets = biobridge_primekg.get_primekg_triplets()
+    biobridge_splits = biobridge_primekg.get_train_test_split()
+    biobridge_node_info = biobridge_primekg.get_node_info_dict()
+
+    # Check if the local directories exists
+    assert os.path.exists(biobridge_primekg.primekg_dir)
+    assert os.path.exists(biobridge_primekg.local_dir)
+    # Check if downloaded and processed files exist
+    # PrimeKG files
+    files = ["nodes.tab", "primekg_nodes.tsv.gz",
+             "edges.csv", "primekg_edges.tsv.gz"]
+    for file in files:
+        path = f"{biobridge_primekg.primekg_dir}/{file}"
+        assert os.path.exists(path)
+    # BioBridge data config
+    assert os.path.exists(f"{biobridge_primekg.local_dir}/data_config.json")
+    # BioBridge embeddings
+    files = [
+        "protein.pkl",
+        "mf.pkl",
+        "cc.pkl",
+        "bp.pkl",
+        "drug.pkl",
+        "disease.pkl",
+        "embedding_dict.pkl"
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/embeddings/{file}"
+        assert os.path.exists(path)
+    # BioBridge processed files
+    files = [
+        "protein.csv",
+        "mf.csv",
+        "cc.csv",
+        "bp.csv",
+        "drug.csv",
+        "disease.csv",
+        "triplet_full.tsv.gz",
+        "triplet_full_altered.tsv.gz",
+        "node_train.tsv.gz",
+        "triplet_train.tsv.gz",
+        "node_test.tsv.gz",
+        "triplet_test.tsv.gz",
+    ]
+    for file in files:
+        path = f"{biobridge_primekg.local_dir}/processed/{file}"
+        assert os.path.exists(path)
+    # Check processed PrimeKG dataframes
+    # Nodes
+    assert primekg_nodes is not None
+    assert len(primekg_nodes) > 0
+    assert primekg_nodes.shape[0] == 129375
+    # Edges
+    assert primekg_edges is not None
+    assert len(primekg_edges) > 0
+    assert primekg_edges.shape[0] == 8100498
+    # Check processed BioBridge data config
+    assert biobridge_data_config is not None
+    assert len(biobridge_data_config) > 0
+    assert len(biobridge_data_config['node_type']) == 10
+    assert len(biobridge_data_config['relation_type']) == 18
+    assert len(biobridge_data_config['emb_dim']) == 6
+    # Check processed BioBridge embeddings
+    assert biobridge_emb_dict is not None
+    assert len(biobridge_emb_dict) > 0
+    assert len(biobridge_emb_dict) == 85466
+    # Check processed BioBridge triplets
+    assert biobridge_triplets is not None
+    assert len(biobridge_triplets) > 0
+    assert biobridge_triplets.shape[0] == 3904610
+    assert list(biobridge_splits.keys()) == ['train', 'node_train', 'test', 'node_test']
+    assert len(biobridge_splits['train']) == 3510930
+    assert len(biobridge_splits['node_train']) == 76486
+    assert len(biobridge_splits['test']) == 393680
+    assert len(biobridge_splits['node_test']) == 8495
+    # Check node info dictionary
+    assert list(biobridge_node_info.keys()) == ['gene/protein',
+                                                'molecular_function',
+                                                'cellular_component',
+                                                'biological_process',
+                                                'drug',
+                                                'disease']
+    assert len(biobridge_node_info['gene/protein']) == 19162
+    assert len(biobridge_node_info['molecular_function']) == 10966
+    assert len(biobridge_node_info['cellular_component']) == 4013
+    assert len(biobridge_node_info['biological_process']) == 27478
+    assert len(biobridge_node_info['drug']) == 6948
+    assert len(biobridge_node_info['disease']) == 44133
+
+# def test_load_existing_primekg_with_negative_triplets(biobridge_primekg):
+#     """
+#     Test the loading method of the BioBridge-PrimeKG class by loading existing data in local.
+#     In addition, it builds negative triplets for training data.
+#     """
+#     # Load BioBridge-PrimeKG data
+#     # Using 1 negative sample per positive triplet
+#     biobridge_primekg.load_data(build_neg_triplest=True, n_neg_samples=1)
+#     biobridge_neg_triplets = biobridge_primekg.get_primekg_triplets_negative()
+
+#     # Check if the local directories exists
+#     assert os.path.exists(biobridge_primekg.primekg_dir)
+#     assert os.path.exists(biobridge_primekg.local_dir)
+#     # Check if downloaded and processed files exist
+#     path = f"{biobridge_primekg.local_dir}/processed/triplet_train_negative.tsv.gz"
+#     assert os.path.exists(path)
+#     # Check processed BioBridge triplets
+#     assert biobridge_neg_triplets is not None
+#     assert len(biobridge_neg_triplets) > 0
+#     assert biobridge_neg_triplets.shape[0] == 3510930
+#     assert len(biobridge_neg_triplets.negative_tail_index[0]) == 1

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_dataset.py

@@ -0,0 +1,29 @@
+"""
+Test cases for datasets/dataset.py
+"""
+
+from ..datasets.dataset import Dataset
+
+class MockDataset(Dataset):
+    """
+    Mock dataset class for testing purposes.
+    """
+    def setup(self):
+        pass
+
+    def load_data(self):
+        pass
+
+def test_dataset_setup():
+    """
+    Test the setup method of the Dataset class.
+    """
+    dataset = MockDataset()
+    assert dataset.setup() is None
+
+def test_dataset_load_data():
+    """
+    Test the load_data method of the Dataset class.
+    """
+    dataset = MockDataset()
+    assert dataset.load_data() is None

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_primekg.py

@@ -0,0 +1,73 @@
+"""
+Test cases for datasets/primekg_loader.py
+"""
+
+import os
+import shutil
+import pytest
+from ..datasets.primekg import PrimeKG
+
+# Remove the data folder for testing if it exists
+LOCAL_DIR = "../data/primekg_test/"
+shutil.rmtree(LOCAL_DIR, ignore_errors=True)
+
+@pytest.fixture(name="primekg")
+def primekg_fixture():
+    """
+    Fixture for creating an instance of PrimeKG.
+    """
+    return PrimeKG(local_dir=LOCAL_DIR)
+
+def test_download_primekg(primekg):
+    """
+    Test the loading method of the PrimeKG class by downloading PrimeKG from server.
+    """
+    # Load PrimeKG data
+    primekg.load_data()
+    primekg_nodes = primekg.get_nodes()
+    primekg_edges = primekg.get_edges()
+
+    # Check if the local directory exists
+    assert os.path.exists(primekg.local_dir)
+    # Check if downloaded and processed files exist
+    files = ["nodes.tab", f"{primekg.name}_nodes.tsv.gz",
+             "edges.csv", f"{primekg.name}_edges.tsv.gz"]
+    for file in files:
+        path = f"{primekg.local_dir}/{file}"
+        assert os.path.exists(path)
+    # Check processed PrimeKG dataframes
+    # Nodes
+    assert primekg_nodes is not None
+    assert len(primekg_nodes) > 0
+    assert primekg_nodes.shape[0] == 129375
+    # Edges
+    assert primekg_edges is not None
+    assert len(primekg_edges) > 0
+    assert primekg_edges.shape[0] == 8100498
+
+def test_load_existing_primekg(primekg):
+    """
+    Test the loading method of the PrimeKG class by loading existing PrimeKG in local.
+    """
+    # Load PrimeKG data
+    primekg.load_data()
+    primekg_nodes = primekg.get_nodes()
+    primekg_edges = primekg.get_edges()
+
+    # Check if the local directory exists
+    assert os.path.exists(primekg.local_dir)
+    # Check if downloaded and processed files exist
+    files = ["nodes.tab", f"{primekg.name}_nodes.tsv.gz",
+             "edges.csv", f"{primekg.name}_edges.tsv.gz"]
+    for file in files:
+        path = f"{primekg.local_dir}/{file}"
+        assert os.path.exists(path)
+    # Check processed PrimeKG dataframes
+    # Nodes
+    assert primekg_nodes is not None
+    assert len(primekg_nodes) > 0
+    assert primekg_nodes.shape[0] == 129375
+    # Edges
+    assert primekg_edges is not None
+    assert len(primekg_edges) > 0
+    assert primekg_edges.shape[0] == 8100498

aiagents4pharma/talk2knowledgegraphs/tests/test_datasets_starkqa_primekg.py

@@ -0,0 +1,116 @@
+"""
+Test cases for datasets/starkqa_primekg_loader.py
+"""
+
+import os
+import shutil
+import pytest
+from ..datasets.starkqa_primekg import StarkQAPrimeKG
+
+# Remove the data folder for testing if it exists
+LOCAL_DIR = "../data/starkqa_primekg_test/"
+shutil.rmtree(LOCAL_DIR, ignore_errors=True)
+
+@pytest.fixture(name="starkqa_primekg")
+def starkqa_primekg_fixture():
+    """
+    Fixture for creating an instance of StarkQAPrimeKGData.
+    """
+    return StarkQAPrimeKG(local_dir=LOCAL_DIR)
+
+def test_download_starkqa_primekg(starkqa_primekg):
+    """
+    Test the loading method of the StarkQAPrimeKGLoaderTool class by downloading files
+    from HuggingFace Hub.
+    """
+    # Load StarkQA PrimeKG data
+    starkqa_primekg.load_data()
+    starkqa_df = starkqa_primekg.get_starkqa()
+    primekg_node_info = starkqa_primekg.get_starkqa_node_info()
+    split_idx = starkqa_primekg.get_starkqa_split_indicies()
+    query_embeddings = starkqa_primekg.get_query_embeddings()
+    node_embeddings = starkqa_primekg.get_node_embeddings()
+
+    # Check if the local directory exists
+    assert os.path.exists(starkqa_primekg.local_dir)
+    # Check if downloaded files exist in the local directory
+    files = ['qa/prime/split/test-0.1.index',
+             'qa/prime/split/test.index',
+             'qa/prime/split/train.index',
+             'qa/prime/split/val.index',
+             'qa/prime/stark_qa/stark_qa.csv',
+             'qa/prime/stark_qa/stark_qa_human_generated_eval.csv',
+             'skb/prime/processed.zip']
+    for file in files:
+        path = f"{starkqa_primekg.local_dir}/{file}"
+        assert os.path.exists(path)
+    # Check dataframe
+    assert starkqa_df is not None
+    assert len(starkqa_df) > 0
+    assert starkqa_df.shape[0] == 11204
+    # Check node information
+    assert primekg_node_info is not None
+    assert len(primekg_node_info) == 129375
+    # Check split indices
+    assert list(split_idx.keys()) == ['train', 'val', 'test', 'test-0.1']
+    assert len(split_idx['train']) == 6162
+    assert len(split_idx['val']) == 2241
+    assert len(split_idx['test']) == 2801
+    assert len(split_idx['test-0.1']) == 280
+    # Check query embeddings
+    assert query_embeddings is not None
+    assert len(query_embeddings) == 11204
+    assert query_embeddings[0].shape[1] == 1536
+    # Check node embeddings
+    assert node_embeddings is not None
+    assert len(node_embeddings) == 129375
+    assert node_embeddings[0].shape[1] == 1536
+
+def test_load_existing_starkqa_primekg(starkqa_primekg):
+    """
+
+    Test the loading method of the StarkQAPrimeKGLoaderTool class by loading existing files
+    in the local directory.
+    """
+    # Load StarkQA PrimeKG data
+    starkqa_primekg.load_data()
+    starkqa_df = starkqa_primekg.get_starkqa()
+    primekg_node_info = starkqa_primekg.get_starkqa_node_info()
+    split_idx = starkqa_primekg.get_starkqa_split_indicies()
+    query_embeddings = starkqa_primekg.get_query_embeddings()
+    node_embeddings = starkqa_primekg.get_node_embeddings()
+
+    # Check if the local directory exists
+    assert os.path.exists(starkqa_primekg.local_dir)
+    # Check if downloaded and processed files exist
+    files = ['qa/prime/split/test-0.1.index',
+             'qa/prime/split/test.index',
+             'qa/prime/split/train.index',
+             'qa/prime/split/val.index',
+             'qa/prime/stark_qa/stark_qa.csv',
+             'qa/prime/stark_qa/stark_qa_human_generated_eval.csv',
+             'skb/prime/processed.zip']
+    for file in files:
+        path = f"{starkqa_primekg.local_dir}/{file}"
+        assert os.path.exists(path)
+    # Check dataframe
+    assert starkqa_df is not None
+    assert len(starkqa_df) > 0
+    assert starkqa_df.shape[0] == 11204
+    # Check node information
+    assert primekg_node_info is not None
+    assert len(primekg_node_info) == 129375
+    # Check split indices
+    assert list(split_idx.keys()) == ['train', 'val', 'test', 'test-0.1']
+    assert len(split_idx['train']) == 6162
+    assert len(split_idx['val']) == 2241
+    assert len(split_idx['test']) == 2801
+    assert len(split_idx['test-0.1']) == 280
+    # Check query embeddings
+    assert query_embeddings is not None
+    assert len(query_embeddings) == 11204
+    assert query_embeddings[0].shape[1] == 1536
+    # Check node embeddings
+    assert node_embeddings is not None
+    assert len(node_embeddings) == 129375
+    assert node_embeddings[0].shape[1] == 1536

aiagents4pharma/talk2knowledgegraphs/tests/test_utils_embeddings_embeddings.py

@@ -0,0 +1,47 @@
+"""
+Test cases for utils/embeddings/embeddings.py
+"""
+
+import pytest
+from ..utils.embeddings.embeddings import Embeddings
+
+class TestEmbeddings(Embeddings):
+    """Test implementation of the Embeddings interface for testing purposes."""
+
+    def embed_documents(self, texts: list[str]) -> list[list[float]]:
+        return [[0.1, 0.2, 0.3] for _ in texts]
+
+    def embed_query(self, text: str) -> list[float]:
+        return [0.1, 0.2, 0.3]
+
+def test_embed_documents():
+    """Test embedding documents using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    texts = ["text1", "text2"]
+    result = embeddings.embed_documents(texts)
+    assert result == [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]
+
+
+def test_embed_query():
+    """Test embedding a query using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    text = "query"
+    result = embeddings.embed_query(text)
+    assert result == [0.1, 0.2, 0.3]
+
+@pytest.mark.asyncio
+async def test_aembed_documents():
+    """Test asynchronous embedding of documents using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    texts = ["text1", "text2"]
+    result = await embeddings.aembed_documents(texts)
+    assert result == [[0.1, 0.2, 0.3], [0.1, 0.2, 0.3]]
+
+
+@pytest.mark.asyncio
+async def test_aembed_query():
+    """Test asynchronous embedding of a query using the Embeddings interface."""
+    embeddings = TestEmbeddings()
+    text = "query"
+    result = await embeddings.aembed_query(text)
+    assert result == [0.1, 0.2, 0.3]