linkml-store 0.3.0__py3-none-any.whl

linkml_store/constants.py

@@ -0,0 +1,7 @@
+import pystow
+
+__all__ = [
+    "LINKML_STORE_MODULE",
+]
+
+LINKML_STORE_MODULE = pystow.module("linkml", "store")

linkml_store/graphs/graph_map.py

@@ -0,0 +1,24 @@
+from abc import ABC
+from typing import Optional
+
+from pydantic import BaseModel
+
+DEFAULT_IDENTIFIER_ATTRIBUTE = "id"
+DEFAULT_CATEGORY_LABELS_ATTRIBUTE = "category"
+DEFAULT_SUBJECT_ATTRIBUTE = "subject"
+DEFAULT_PREDICATE_ATTRIBUTE = "predicate"
+DEFAULT_OBJECT_ATTRIBUTE = "object"
+
+
+class GraphProjection(BaseModel, ABC):
+    identifier_attribute: str = DEFAULT_IDENTIFIER_ATTRIBUTE
+
+
+class NodeProjection(GraphProjection):
+    category_labels_attribute: Optional[str] = DEFAULT_CATEGORY_LABELS_ATTRIBUTE
+
+
+class EdgeProjection(GraphProjection):
+    subject_attribute: str = DEFAULT_SUBJECT_ATTRIBUTE
+    predicate_attribute: str = DEFAULT_PREDICATE_ATTRIBUTE
+    object_attribute: str = DEFAULT_OBJECT_ATTRIBUTE

linkml_store/index/__init__.py

@@ -0,0 +1,53 @@
+"""
+Indexers package.
+
+Indexers allow indexes to be added to existing :class:`Collection` objects.
+
+Current two are supported:
+
+* simple: :class:`SimpleIndexer`
+* llm: :class:`LLMIndexer`
+"""
+
+from typing import Type
+
+from linkml_store.index.implementations.llm_indexer import LLMIndexer
+from linkml_store.index.implementations.simple_indexer import SimpleIndexer
+from linkml_store.index.indexer import Indexer
+
+INDEXER_CLASSES = {
+    "simple": SimpleIndexer,
+    "llm": LLMIndexer,
+}
+
+
+def get_indexer_class(name: str) -> Type[Indexer]:
+    """
+    Get an indexer class by name.
+
+    :param name: the name of the indexer (simple, llm, ...)
+    :return: the indexer class
+    """
+    if name not in INDEXER_CLASSES:
+        raise ValueError(f"Unknown indexer class: {name}")
+    return INDEXER_CLASSES[name]
+
+
+def get_indexer(index_type: str, **kwargs) -> Indexer:
+    """
+    Get an indexer by name.
+
+    >>> simple_indexer = get_indexer("simple")
+    >>> llm_indexer = get_indexer("llm")
+
+    :param name: the name of the indexer (simple, llm, ...)
+    :param kwargs: additional arguments to pass to the indexer
+    :return: the indexer
+    """
+    kwargs = {k: v for k, v in kwargs.items() if v is not None}
+    cls = get_indexer_class(index_type)
+    kwargs["index_type"] = index_type
+    indexer = cls(**kwargs)
+    if not indexer.name:
+        indexer.name = index_type
+    return indexer

linkml_store/index/implementations/llm_indexer.py

@@ -0,0 +1,174 @@
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, List, Optional
+
+import numpy as np
+
+from linkml_store.api.config import CollectionConfig
+from linkml_store.index.indexer import INDEX_ITEM, Indexer
+from linkml_store.utils.llm_utils import get_token_limit, render_formatted_text
+
+if TYPE_CHECKING:
+    import llm
+
+CHUNK_SIZE = 1000
+
+logger = logging.getLogger(__name__)
+
+
+class LLMIndexer(Indexer):
+    """
+    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")
+
+    TODO: Implement true batching for embedding API calls
+    TODO: Add batch_size parameter to control batch processing
+    TODO: Support batch embedding APIs (e.g., OpenAI batch endpoint)
+    TODO: Add progress reporting for large batch operations
+    TODO: Implement smart batching with accumulation and flushing
+    """
+
+    embedding_model_name: str = "text-embedding-ada-002"
+    _embedding_model: "llm.EmbeddingModel" = None
+    cached_embeddings_database: str = None
+    cached_embeddings_collection: str = None
+    cache_queries: bool = False
+    truncation_method: Optional[str] = None
+    # TODO: Add batch_size: int = 100 parameter for batch processing
+    # TODO: Add supported_models class variable with model metadata (dims, costs, limits)
+    # TODO: Add model_validation to check if model exists before use
+
+    @property
+    def embedding_model(self):
+        import llm
+
+        if self._embedding_model is None:
+            self._embedding_model = llm.get_embedding_model(self.embedding_model_name)
+        return self._embedding_model
+
+    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], cache=cache, **kwargs)[0]
+
+    def texts_to_vectors(
+        self, texts: List[str], cache: bool = None, token_limit_penalty=0, batch_size: int=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:
+        :param cache:
+        :param token_limit_penalty:
+        :return:
+        """
+        from tiktoken import encoding_for_model
+
+        logging.info(f"Converting {len(texts)} texts to vectors")
+        model = self.embedding_model
+        # TODO: make this more accurate
+        token_limit = get_token_limit(model.model_id) - token_limit_penalty
+        logging.info(f"Token limit for {model.model_id}: {token_limit}")
+        encoding = encoding_for_model(self.embedding_model_name)
+
+        def truncate_text(text: str) -> str:
+            # split into tokens every 1000 chars:
+            parts = [text[i : i + CHUNK_SIZE] for i in range(0, len(text), CHUNK_SIZE)]
+            truncated = render_formatted_text(
+                lambda x: "".join(x),
+                parts,
+                encoding,
+                token_limit,
+            )
+            logger.debug(f"Truncated text from {len(text)} to {len(truncated)}")
+            return truncated
+
+        texts = [truncate_text(text) for text in texts]
+        # Calculate average number of tokens per text for accurate batch sizing
+        text_token_counts = [len(encoding.encode(t)) for t in texts]
+        avg_text_tokens = sum(text_token_counts) / len(text_token_counts)
+        logger.info(f"Average text token count: {avg_text_tokens}")
+        if batch_size is None:
+            # TODO: empirically determine best batch size
+            batch_size = max(int(token_limit / avg_text_tokens), 5)
+            logger.info(f"Setting batch size to {batch_size}")
+        
+
+        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(
+                alias=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)
+
+            embeddings = list([None] * len(texts))
+            uncached_texts = []
+            n = 0
+            # TODO: Implement batch lookup for cache checking (single query for all texts)
+            # TODO: Use IN clause or batch query to check multiple texts at once
+            logger.info(f"Checking cache for {len(texts)} texts")
+            for i in range(len(texts)):
+                # TODO: optimize this - currently makes N database queries for N texts
+                text = texts[i]
+                logger.debug(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, batch_size=batch_size))
+                # TODO: Combine into a single insert with multiple rows for better performance
+                # TODO: Use insert_many or bulk insert instead of individual inserts
+                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}
+                    )
+                embeddings_collection.commit()
+        else:
+            logger.info(f"Embedding {len(texts)} texts")
+            # TODO: Add progress callback for large batches without cache
+            embeddings = list(model.embed_multi(texts, batch_size=batch_size))
+        return [np.array(v, dtype=float) for v in embeddings]

linkml_store/index/implementations/simple_indexer.py

@@ -0,0 +1,43 @@
+import hashlib
+import logging
+
+import numpy as np
+
+from linkml_store.index.indexer import INDEX_ITEM, Indexer
+
+logger = logging.getLogger(__name__)
+
+
+class SimpleIndexer(Indexer):
+    """
+    A implementations index that uses a hash function to generate an index from text.
+
+    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, cache: bool = None, **kwargs) -> INDEX_ITEM:
+        """
+        This is a naive method purely for testing
+
+        :param text:
+        :return:
+        """
+        vector_length = self.vector_default_length
+        text = text.lower()
+        # trigrams
+        words = [text[i : i + 3] for i in range(len(text) - 2)]
+
+        vector = np.zeros(vector_length, dtype=float)
+
+        # Iterate over each trigram in the text
+        for word in words:
+            # Generate a hash value for the word
+            hash_value = int(hashlib.sha1(word.encode("utf-8")).hexdigest(), 16)
+
+            # Compute the index in the vector using modulo
+            index = hash_value % vector_length
+
+            # Increment the count at the computed index
+            vector[index] += 1.0
+        logger.debug(f"Indexed text: {text} as {vector}")
+        return vector

linkml_store/index/indexer.py

@@ -0,0 +1,211 @@
+import logging
+from enum import Enum
+from typing import Any, Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+from pydantic import BaseModel
+
+from linkml_store.utils.vector_utils import mmr_diversified_search, pairwise_cosine_similarity
+
+INDEX_ITEM = np.ndarray
+
+logger = logging.getLogger(__name__)
+
+
+class TemplateSyntaxEnum(str, Enum):
+    """
+    Template syntax types.
+    """
+
+    jinja2 = "jinja2"
+    fstring = "fstring"
+
+
+class Indexer(BaseModel):
+    """
+    An indexer operates on a collection in order to search for objects.
+
+    You should use a subcllass of this; this can be looked up dynqamically:
+
+    >>> from linkml_store.index import get_indexer
+    >>> indexer = get_indexer("simple")
+
+    You can customize how objects are indexed by passing in a text template.
+    For example, if your collection has objects with "name" and "profession" attributes,
+    you can index them as "{name} {profession}".
+
+    >>> indexer = get_indexer("simple", text_template="{name} :: {profession}")
+
+    By default, python fstrings are assumed.
+
+    We can test this works using the :ref:`object_to_text` method (normally
+    you would never need to call this directly, but it's useful for testing):
+
+    >>> obj = {"name": "John", "profession": "doctor"}
+    >>> indexer.object_to_text(obj)
+    'John :: doctor'
+
+    You can also use Jinja2 templates; this gives more flexibility and logic,
+    e.g. conditional formatting:
+
+    >>> tmpl = "{{name}}{% if profession %} :: {{profession}}{% endif %}"
+    >>> indexer = get_indexer("simple", text_template=tmpl, text_template_syntax=TemplateSyntaxEnum.jinja2)
+    >>> indexer.object_to_text(obj)
+    'John :: doctor'
+    >>> indexer.object_to_text({"name": "John"})
+    'John'
+
+    You can also specify which attributes to index:
+
+    >>> indexer = get_indexer("simple", index_attributes=["name"])
+    >>> indexer.object_to_text(obj)
+    'John'
+
+    The purpose of an indexer is to translate a collection of objects into a collection of objects
+    such as vectors for purposes such as search. Unless you are implementing your own indexer, you
+    generally don't need to use the methods that return vectors, but we can examine their behavior
+    to get a sense of how they work.
+
+    >>> vectors = indexer.objects_to_vectors([{"name": "Aardvark"}, {"name": "Aardwolf"}, {"name": "Zesty"}])
+    >>> assert pairwise_cosine_similarity(vectors[0], vectors[1]) > pairwise_cosine_similarity(vectors[0], vectors[2])
+
+    Note you should consult the documentation for the specific indexer you are using for more details on
+    how text is converted to vectors.
+
+    """
+
+    name: Optional[str] = None
+    index_type: Optional[str] = None
+    index_function: Optional[Callable] = None
+    distance_function: Optional[Callable] = 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__"
+    index_value_field: Optional[str] = "__index_value__"
+
+    def object_to_vector(self, obj: Dict[str, Any]) -> INDEX_ITEM:
+        """
+        Convert an object to an indexable object
+
+        :param obj:
+        :return:
+        """
+        return self.text_to_vector(self.object_to_text(obj))
+
+    def objects_to_vectors(self, objs: List[Dict[str, Any]]) -> List[INDEX_ITEM]:
+        """
+        Convert a list of objects to indexable objects
+
+        :param objs:
+        :return: list of vectors
+        """
+        return self.texts_to_vectors([self.object_to_text(obj) for obj in objs])
+
+    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, cache=cache, **kwargs) for text in texts]
+
+    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
+
+    def object_to_text(self, obj: Dict[str, Any]) -> str:
+        """
+        Convert an object to a text representation
+
+        :param obj:
+        :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:
+            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 not syntax:
+                syntax = TemplateSyntaxEnum.fstring
+            if syntax == TemplateSyntaxEnum.jinja2:
+                from jinja2 import Template
+
+                template = Template(self.text_template)
+                return template.render(**obj)
+            elif syntax == TemplateSyntaxEnum.fstring:
+                return self.text_template.format(**obj)
+            else:
+                raise NotImplementedError(f"Cannot handle template syntax: {syntax}")
+        return str(obj)
+
+    def search(
+        self,
+        query: str,
+        vectors: List[Tuple[str, INDEX_ITEM]],
+        limit: Optional[int] = None,
+        mmr_relevance_factor: Optional[float] = None,
+    ) -> List[Tuple[float, Any]]:
+        """
+        Use the indexer to search against a database of vectors.
+
+        Note: this is a low-level method, typically you would use the :ref:`search` method on a :ref:`Collection`.
+
+        :param query: The query string to search for
+        :param vectors: A list of indexed items, where each item is a tuple of (id, vector)
+        :param limit: The maximum number of results to return (optional)
+        :return: A list of item IDs or objects that match the query
+        """
+
+        # Convert the query string to a vector
+        query_vector = self.text_to_vector(query, cache=False)
+
+        if mmr_relevance_factor is not None:
+            vlist = [v for _, v in vectors]
+            idlist = [id for id, _ in vectors]
+            sorted_indices = mmr_diversified_search(
+                query_vector, vlist, relevance_factor=mmr_relevance_factor, top_n=limit
+            )
+            results = []
+            # TODO: this is inefficient when limit is high
+            for i in range(limit):
+                if i >= len(sorted_indices):
+                    break
+                pos = sorted_indices[i]
+                score = pairwise_cosine_similarity(query_vector, vlist[pos])
+                results.append((score, idlist[pos]))
+            return results
+
+        distances = []
+
+        # Iterate over each indexed item
+        for item_id, item_vector in vectors:
+            # Calculate the Euclidean distance between the query vector and the item vector
+            # distance = 1-np.linalg.norm(query_vector - item_vector)
+            distance = pairwise_cosine_similarity(query_vector, item_vector)
+            distances.append((distance, item_id))
+
+        # Sort the distances in ascending order
+        distances.sort(key=lambda x: -x[0])
+
+        # Limit the number of results if specified
+        if limit is not None:
+            distances = distances[:limit]
+
+        return distances

linkml_store/inference/__init__.py

@@ -0,0 +1,13 @@
+"""
+inference engine package.
+"""
+
+from linkml_store.inference.inference_config import InferenceConfig
+from linkml_store.inference.inference_engine import InferenceEngine
+from linkml_store.inference.inference_engine_registry import get_inference_engine
+
+__all__ = [
+    "InferenceEngine",
+    "InferenceConfig",
+    "get_inference_engine",
+]