linkml-store 0.1.7__py3-none-any.whl → 0.1.9__py3-none-any.whl

linkml_store/index/implementations/llm_indexer.py

@@ -1,20 +1,34 @@
+import logging
+from pathlib import Path
 from typing import TYPE_CHECKING, List
 
 import numpy as np
 
+from linkml_store.api.config import CollectionConfig
 from linkml_store.index.indexer import INDEX_ITEM, Indexer
 
 if TYPE_CHECKING:
     import llm
 
 
+logger = logging.getLogger(__name__)
+
+
 class LLMIndexer(Indexer):
     """
-    A implementations index wraps the llm library
+    An indexer that wraps the llm library.
+
+    This indexer is used to convert text to vectors using the llm library.
+
+    >>> indexer = LLMIndexer(cached_embeddings_database="tests/input/llm_cache.db")
+    >>> vector = indexer.text_to_vector("hello")
     """
 
     embedding_model_name: str = "ada-002"
     _embedding_model: "llm.EmbeddingModel" = None
+    cached_embeddings_database: str = None
+    cached_embeddings_collection: str = None
+    cache_queries: bool = False
 
     @property
     def embedding_model(self):
@@ -24,21 +38,85 @@ class LLMIndexer(Indexer):
             self._embedding_model = llm.get_embedding_model(self.embedding_model_name)
         return self._embedding_model
 
-    def text_to_vector(self, text: str) -> INDEX_ITEM:
+    def text_to_vector(self, text: str, cache: bool = None, **kwargs) -> INDEX_ITEM:
         """
         Convert a text to an indexable object
 
+        >>> indexer = LLMIndexer(cached_embeddings_database="tests/input/llm_cache.db")
+        >>> vector = indexer.text_to_vector("hello")
+
         :param text:
         :return:
         """
-        return self.texts_to_vectors([text])[0]
+        return self.texts_to_vectors([text], cache=cache, **kwargs)[0]
 
-    def texts_to_vectors(self, texts: List[str]) -> List[INDEX_ITEM]:
+    def texts_to_vectors(self, texts: List[str], cache: bool = None, **kwargs) -> List[INDEX_ITEM]:
         """
         Use LLM to embed
 
+        >>> indexer = LLMIndexer(cached_embeddings_database="tests/input/llm_cache.db")
+        >>> vectors = indexer.texts_to_vectors(["hello", "goodbye"])
+
         :param texts:
         :return:
         """
-        embeddings = self.embedding_model.embed_multi(texts)
+        logging.info(f"Converting {len(texts)} texts to vectors")
+        model = self.embedding_model
+        if self.cached_embeddings_database and (cache is None or cache or self.cache_queries):
+            model_id = model.model_id
+            if not model_id:
+                raise ValueError("Model ID is required to cache embeddings")
+            db_path = Path(self.cached_embeddings_database)
+            coll_name = self.cached_embeddings_collection
+            if not coll_name:
+                coll_name = "all_embeddings"
+            from linkml_store import Client
+
+            embeddings_client = Client()
+            config = CollectionConfig(
+                name=coll_name,
+                type="Embeddings",
+                attributes={
+                    "text": {"range": "string"},
+                    "model_id": {"range": "string"},
+                    "embedding": {"range": "float", "array": {}},
+                },
+            )
+            embeddings_db = embeddings_client.get_database(f"duckdb:///{db_path}")
+            if coll_name in embeddings_db.list_collection_names():
+                # Load existing collection and use its model
+                embeddings_collection = embeddings_db.create_collection(coll_name, metadata=config)
+            else:
+                embeddings_collection = embeddings_db.create_collection(coll_name, metadata=config)
+            texts = list(texts)
+            embeddings = list([None] * len(texts))
+            uncached_texts = []
+            n = 0
+            for i in range(len(texts)):
+                # TODO: optimize this
+                text = texts[i]
+                logger.info(f"Looking for cached embedding for {text}")
+                r = embeddings_collection.find({"text": text, "model_id": model_id})
+                if r.num_rows:
+                    embeddings[i] = r.rows[0]["embedding"]
+                    n += 1
+                    logger.info("Found")
+                else:
+                    uncached_texts.append((text, i))
+                    logger.info("NOT Found")
+            logger.info(f"Found {n} cached embeddings")
+            if uncached_texts:
+                logger.info(f"Embedding {len(uncached_texts)} uncached texts")
+                uncached_texts, uncached_indices = zip(*uncached_texts)
+                uncached_embeddings = list(model.embed_multi(uncached_texts))
+                # TODO: combine into a single insert with multiple rows
+                for i, index in enumerate(uncached_indices):
+                    logger.debug(f"Indexing text at {i}")
+                    embeddings[index] = uncached_embeddings[i]
+                    embeddings_collection.insert(
+                        {"text": uncached_texts[i], "embedding": embeddings[index], "model_id": model_id}
+                    )
+        else:
+            logger.info(f"Embedding {len(texts)} texts")
+            embeddings = model.embed_multi(texts)
         return [np.array(v, dtype=float) for v in embeddings]

linkml_store/index/implementations/simple_indexer.py

@@ -15,7 +15,7 @@ class SimpleIndexer(Indexer):
     This uses a naive method to generate an index from text. It is not suitable for production use.
     """
 
-    def text_to_vector(self, text: str) -> INDEX_ITEM:
+    def text_to_vector(self, text: str, cache: bool = None, **kwargs) -> INDEX_ITEM:
         """
         This is a naive method purely for testing
 
@@ -39,5 +39,5 @@ class SimpleIndexer(Indexer):
 
             # Increment the count at the computed index
             vector[index] += 1.0
-        logger.info(f"Indexed text: {text} as {vector}")
+        logger.debug(f"Indexed text: {text} as {vector}")
         return vector

linkml_store/index/indexer.py

@@ -1,3 +1,5 @@
+import logging
+from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import numpy as np
@@ -5,6 +7,13 @@ from pydantic import BaseModel
 
 INDEX_ITEM = np.ndarray
 
+logger = logging.getLogger(__name__)
+
+
+class TemplateSyntaxEnum(str, Enum):
+    jinja2 = "jinja2"
+    fstring = "fstring"
+
 
 def cosine_similarity(vector1, vector2):
     dot_product = np.dot(vector1, vector2)
@@ -21,8 +30,9 @@ class Indexer(BaseModel):
     name: Optional[str] = None
     index_function: Optional[Callable] = None
     distance_function: Optional[Callable] = None
-    index_attributes: Optional[str] = None
+    index_attributes: Optional[List[str]] = None
     text_template: Optional[str] = None
+    text_template_syntax: Optional[TemplateSyntaxEnum] = None
     filter_nulls: Optional[bool] = True
     vector_default_length: Optional[int] = 1000
     index_field: Optional[str] = "__index__"
@@ -41,24 +51,25 @@ class Indexer(BaseModel):
         Convert a list of objects to indexable objects
 
         :param objs:
-        :return:
+        :return: list of vectors
         """
-        return [self.object_to_vector(obj) for obj in objs]
+        return self.texts_to_vectors([self.object_to_text(obj) for obj in objs])
 
-    def texts_to_vectors(self, texts: List[str]) -> List[INDEX_ITEM]:
+    def texts_to_vectors(self, texts: List[str], cache: bool = None, **kwargs) -> List[INDEX_ITEM]:
         """
         Convert a list of texts to indexable objects
 
         :param texts:
         :return:
         """
-        return [self.text_to_vector(text) for text in texts]
+        return [self.text_to_vector(text, cache=cache, **kwargs) for text in texts]
 
-    def text_to_vector(self, text: str) -> INDEX_ITEM:
+    def text_to_vector(self, text: str, cache: bool = None, **kwargs) -> INDEX_ITEM:
         """
         Convert a text to an indexable object
 
         :param text:
+        :param cache:
         :return:
         """
         raise NotImplementedError
@@ -71,11 +82,24 @@ class Indexer(BaseModel):
         :return:
         """
         if self.index_attributes:
+            if len(self.index_attributes) == 1 and not self.text_template:
+                return str(obj[self.index_attributes[0]])
             obj = {k: v for k, v in obj.items() if k in self.index_attributes}
         if self.filter_nulls:
             obj = {k: v for k, v in obj.items() if v is not None}
         if self.text_template:
-            return self.text_template.format(**obj)
+            syntax = self.text_template_syntax
+            if not syntax:
+                if "{%" in self.text_template or "{{" in self.text_template:
+                    logger.info("Detected Jinja2 syntax in text template")
+                    syntax = TemplateSyntaxEnum.jinja2
+            if syntax and syntax == TemplateSyntaxEnum.jinja2:
+                from jinja2 import Template
+
+                template = Template(self.text_template)
+                return template.render(**obj)
+            else:
+                return self.text_template.format(**obj)
         return str(obj)
 
     def search(
@@ -91,7 +115,7 @@ class Indexer(BaseModel):
         """
 
         # Convert the query string to a vector
-        query_vector = self.text_to_vector(query)
+        query_vector = self.text_to_vector(query, cache=False)
 
         distances = []
 

linkml_store/utils/change_utils.py

@@ -0,0 +1,17 @@
+from typing import List
+
+from linkml_store.api.collection import OBJECT
+
+
+def insert_operation_to_patches(objs: List[OBJECT], **kwargs):
+    """
+    Translate a list of objects to a list of patches for insertion.
+
+    Note: inserts are always treated as being at the start of a list
+
+    :param objs: objects to insert
+    :param kwargs: additional arguments
+    """
+    patches = []
+    for obj in objs:
+        patches.append({"op": "add", "path": "/0", "value": obj})

linkml_store/utils/format_utils.py

@@ -4,26 +4,40 @@ import sys
 from enum import Enum
 from io import StringIO
 from pathlib import Path
-from typing import Any, Dict, List, Union
+from typing import Any, Dict, List, Optional, TextIO, Type, Union
 
+import pandas as pd
 import yaml
 from pydantic import BaseModel
 
 
 class Format(Enum):
+    """
+    Supported generic file formats for loading and rendering objects.
+    """
+
     JSON = "json"
     JSONL = "jsonl"
     YAML = "yaml"
     TSV = "tsv"
     CSV = "csv"
+    PARQUET = "parquet"
+    FORMATTED = "formatted"
 
 
-def load_objects(file_path: Union[str, Path], format: Union[Format, str] = None) -> List[Dict[str, Any]]:
+def load_objects(
+    file_path: Union[str, Path], format: Union[Format, str] = None, expected_type: Type = None
+) -> List[Dict[str, Any]]:
     """
     Load objects from a file in JSON, JSONLines, YAML, CSV, or TSV format.
 
+    >>> load_objects("tests/input/test_data/data.csv")
+    [{'id': '1', 'name': 'John', 'age': '30'},
+     {'id': '2', 'name': 'Alice', 'age': '25'}, {'id': '3', 'name': 'Bob', 'age': '35'}]
+
     :param file_path: The path to the file.
     :param format: The format of the file. Can be a Format enum or a string value.
+    :param expected_type: The target type to load the objects into.
     :return: A list of dictionaries representing the loaded objects.
     """
     if isinstance(format, str):
@@ -32,24 +46,39 @@ def load_objects(file_path: Union[str, Path], format: Union[Format, str] = None)
     if isinstance(file_path, Path):
         file_path = str(file_path)
 
+    if not format and (file_path.endswith(".parquet") or file_path.endswith(".pq")):
+        format = Format.PARQUET
+
+    mode = "r"
+    if format == Format.PARQUET:
+        mode = "rb"
+
     if file_path == "-":
         # set file_path to be a stream from stdin
         f = sys.stdin
     else:
-        f = open(file_path)
+        f = open(file_path, mode)
 
     if format == Format.JSON or (not format and file_path.endswith(".json")):
         objs = json.load(f)
     elif format == Format.JSONL or (not format and file_path.endswith(".jsonl")):
         objs = [json.loads(line) for line in f]
     elif format == Format.YAML or (not format and (file_path.endswith(".yaml") or file_path.endswith(".yml"))):
-        objs = yaml.safe_load(f)
+        if expected_type and expected_type == list:
+            objs = list(yaml.safe_load_all(f))
+        else:
+            objs = yaml.safe_load(f)
     elif format == Format.TSV or (not format and file_path.endswith(".tsv")):
         reader = csv.DictReader(f, delimiter="\t")
         objs = list(reader)
     elif format == Format.CSV or (not format and file_path.endswith(".csv")):
         reader = csv.DictReader(f)
         objs = list(reader)
+    elif format == Format.PARQUET:
+        import pyarrow.parquet as pq
+
+        table = pq.read_table(f)
+        objs = table.to_pandas().to_dict(orient="records")
     else:
         raise ValueError(f"Unsupported file format: {file_path}")
     if not isinstance(objs, list):
@@ -57,10 +86,56 @@ def load_objects(file_path: Union[str, Path], format: Union[Format, str] = None)
     return objs
 
 
-def render_output(data: List[Dict[str, Any]], format: Union[Format, str] = Format.YAML) -> str:
+def write_output(
+    data: Union[List[Dict[str, Any]], Dict[str, Any], pd.DataFrame],
+    format: Union[Format, str] = Format.YAML,
+    target: Optional[Union[TextIO, str, Path]] = None,
+) -> None:
+    """
+    Write output data to a file in JSON, JSONLines, YAML, CSV, or TSV format.
+
+    >>> write_output([{"a": 1, "b": 2}, {"a": 3, "b": 4}], Format.JSON, sys.stdout)
+    [
+      {
+        "a": 1,
+        "b": 2
+      },
+      {
+        "a": 3,
+        "b": 4
+        }
+    ]
+    """
+    output_str = render_output(data, format)
+    if target:
+        if isinstance(target, str):
+            with open(target, "w") as target:
+                target.write(output_str)
+        else:
+            target.write(output_str)
+    else:
+        print(output_str)
+
+
+def render_output(
+    data: Union[List[Dict[str, Any]], Dict[str, Any], pd.DataFrame], format: Union[Format, str] = Format.YAML
+) -> str:
     """
     Render output data in JSON, JSONLines, YAML, CSV, or TSV format.
 
+    >>> print(render_output([{"a": 1, "b": 2}, {"a": 3, "b": 4}], Format.JSON))
+    [
+      {
+        "a": 1,
+        "b": 2
+      },
+      {
+        "a": 3,
+        "b": 4
+      }
+    ]
+
+
     :param data: The data to be rendered.
     :param format: The desired output format. Can be a Format enum or a string value.
     :return: The rendered output as a string.
@@ -68,6 +143,14 @@ def render_output(data: List[Dict[str, Any]], format: Union[Format, str] = Forma
     if isinstance(format, str):
         format = Format(format)
 
+    if format == Format.FORMATTED:
+        if not isinstance(data, pd.DataFrame):
+            data = pd.DataFrame(data)
+        return str(data)
+
+    if isinstance(data, pd.DataFrame):
+        data = data.to_dict(orient="records")
+
     if isinstance(data, BaseModel):
         data = data.model_dump()
 
@@ -76,18 +159,66 @@ def render_output(data: List[Dict[str, Any]], format: Union[Format, str] = Forma
     elif format == Format.JSONL:
         return "\n".join(json.dumps(obj) for obj in data)
     elif format == Format.YAML:
-        return yaml.safe_dump(data, sort_keys=False)
+        if isinstance(data, list):
+            return yaml.safe_dump_all(data, sort_keys=False)
+        else:
+            return yaml.safe_dump(data, sort_keys=False)
     elif format == Format.TSV:
         output = StringIO()
-        writer = csv.DictWriter(output, fieldnames=data[0].keys(), delimiter="\t")
+        writer = csv.DictWriter(output, fieldnames=get_fieldnames(data), delimiter="\t")
         writer.writeheader()
         writer.writerows(data)
         return output.getvalue()
     elif format == Format.CSV:
         output = StringIO()
-        writer = csv.DictWriter(output, fieldnames=data[0].keys())
+        writer = csv.DictWriter(output, fieldnames=get_fieldnames(data))
         writer.writeheader()
         writer.writerows(data)
         return output.getvalue()
     else:
         raise ValueError(f"Unsupported output format: {format}")
+
+
+def get_fieldnames(data: List[Dict[str, Any]]) -> List[str]:
+    """
+    Get the fieldnames of a list of dictionaries.
+
+    >>> get_fieldnames([{"a": 1, "b": 2}, {"a": 3, "b": 4}])
+    ['a', 'b']
+
+    :param data: The list of dictionaries.
+    :return: The fieldnames.
+    """
+    fieldnames = []
+    for obj in data:
+        fieldnames.extend([k for k in obj.keys() if k not in fieldnames])
+    return fieldnames
+
+
+def guess_format(path: str) -> Optional[Format]:
+    """
+    Guess the format of a file based on its extension.
+
+    >>> guess_format("data.json")
+    <Format.JSON: 'json'>
+    >>> guess_format("data.jsonl")
+    <Format.JSONL: 'jsonl'>
+    >>> guess_format("data.yaml")
+    <Format.YAML: 'yaml'>
+    >>> assert not guess_format("data")
+
+    :param path: The path to the file.
+    :return: The guessed format.
+    """
+    if path.endswith(".json"):
+        return Format.JSON
+    elif path.endswith(".jsonl"):
+        return Format.JSONL
+    elif path.endswith(".yaml") or path.endswith(".yml"):
+        return Format.YAML
+    elif path.endswith(".tsv"):
+        return Format.TSV
+    elif path.endswith(".csv"):
+        return Format.CSV
+    else:
+        return None

linkml_store/utils/patch_utils.py

@@ -0,0 +1,126 @@
+from typing import Any, Dict, List, Optional, TypedDict
+
+import jsonpatch
+
+
+class PatchDict(TypedDict):
+    op: str
+    path: str
+    value: Optional[Any]
+    _from: Optional[str]
+
+
+def apply_patches(obj: Any, patches: List[PatchDict], primary_key: Optional[str] = None, in_place=False) -> Any:
+    """
+    Apply a set of patches to an object.
+
+    If the object is a list, the primary key must be specified.
+
+    >>> objs = [{'id': 'F1', 'name': 'Cheese'}, {'id': 'F2', 'name': 'Bread'}]
+    >>> patches = [{'op': 'replace', 'path': '/F1/name', 'value': 'Toast'}]
+    >>> apply_patches(objs, patches, primary_key='id')
+    [{'id': 'F1', 'name': 'Toast'}, {'id': 'F2', 'name': 'Bread'}]
+
+    :param obj: object to patch
+    :param patches: list of patches, conforming to the JSON Patch format
+    :param primary_key: key to use as the primary key for the objects (if obj is a list)
+    :param in_place: whether to apply the patches in place
+    :return:
+    """
+    if isinstance(obj, dict):
+        patch_obj = jsonpatch.JsonPatch(patches)
+        return patch_obj.apply(obj, in_place=in_place)
+    elif isinstance(obj, list):
+        if not primary_key:
+            raise ValueError("Primary key must be specified for list objects")
+        return apply_patches_to_list(obj, patches, primary_key, in_place=in_place)
+    else:
+        raise ValueError(f"Unsupported object type: {type(obj)}")
+
+
+def apply_patches_to_list(
+    objects: List[Dict[str, Any]], patches: List[PatchDict], primary_key: str, in_place=False
+) -> List[Dict[str, Any]]:
+    """
+    Apply a set of patches to a list of objects.
+
+
+
+    :param objects: list of objects
+    :param patches: list of patches, conforming to the JSON Patch format
+    :param primary_key: key to use as the primary key for the objects
+    :param in_place: whether to apply the patches in place
+    :return:
+    """
+    objs_as_dict = {obj[primary_key]: obj for obj in objects}
+    result = apply_patches_to_keyed_list(objs_as_dict, patches, in_place=in_place)
+    return list(result.values())
+
+
+def apply_patches_to_keyed_list(
+    objs_as_dict: Dict[str, Dict[str, Any]], patches: List[PatchDict], in_place=False
+) -> Dict[str, Dict[str, Any]]:
+    """
+    Apply a set of patches to a list of objects, where the objects are keyed by a primary key
+
+    :param objs_as_dict:
+    :param patches:
+    :param in_place:
+    :return:
+    """
+    patch_obj = jsonpatch.JsonPatch(patches)
+    result = patch_obj.apply(objs_as_dict, in_place=in_place)
+    return result
+
+
+def patches_from_objects_lists(
+    src_objs: List[Dict[str, Any]], dst_objs: List[Dict[str, Any]], primary_key: str, exclude_none=True
+) -> List[PatchDict]:
+    """
+    Generate a set of patches to transform src_objs into tgt_objs.
+
+    >>> src_objs = [{'id': 'F1', 'name': 'Cheese'}, {'id': 'F2', 'name': 'Bread'}]
+    >>> tgt_objs = [{'id': 'F1', 'name': 'Toast'}, {'id': 'F2', 'name': 'Bread'}]
+    >>> patches_from_objects_lists(src_objs, tgt_objs, primary_key='id')
+    [{'op': 'replace', 'path': '/F1/name', 'value': 'Toast'}]
+
+    by default exclude_none is True, so None values are excluded from the patch
+
+    >>> tgt_objs = [{'id': 'F1', 'name': 'Toast'}, {'id': 'F2', 'name': None}]
+    >>> patches_from_objects_lists(src_objs, tgt_objs, primary_key='id')
+    [{'op': 'replace', 'path': '/F1/name', 'value': 'Toast'}, {'op': 'remove', 'path': '/F2/name'}]
+
+    if exclude_none is False, None values are treated as being set to None
+
+    >>> patches_from_objects_lists(src_objs, tgt_objs, primary_key='id', exclude_none=False)
+    [{'op': 'replace', 'path': '/F1/name', 'value': 'Toast'}, {'op': 'replace', 'path': '/F2/name', 'value': None}]
+
+    See also: `<https://github.com/orgs/linkml/discussions/1975>`_
+
+    Note the patches are sorted deterministically, first by path, then by operation.
+    This helps ensure operations on the same object are grouped together
+
+    :param src_objs: source objects
+    :param dst_objs: target objects
+    :param primary_key: key to use as the primary key for the objects
+    :param exclude_none: whether to exclude None values from the patch
+    :return:
+    """
+    src_objs_as_dict = {obj[primary_key]: obj for obj in src_objs}
+    dst_objs_as_dict = {obj[primary_key]: obj for obj in dst_objs}
+    if exclude_none:
+        src_objs_as_dict = {k: remove_nones(v) for k, v in src_objs_as_dict.items()}
+        dst_objs_as_dict = {k: remove_nones(v) for k, v in dst_objs_as_dict.items()}
+    patch_obj = jsonpatch.JsonPatch.from_diff(src_objs_as_dict, dst_objs_as_dict)
+    pl = patch_obj.patch
+    return sorted(pl, key=lambda x: (x["path"], x["op"]))
+
+
+def remove_nones(obj: Dict[str, Any]) -> Dict[str, Any]:
+    """
+    Remove None values from a dictionary.
+
+    :param obj:
+    :return:
+    """
+    return {k: v for k, v in obj.items() if v is not None}