linkml-store 0.1.12__py3-none-any.whl → 0.1.14__py3-none-any.whl

linkml_store/cli.py

@@ -1,6 +1,7 @@
 import logging
 import sys
 import warnings
+from pathlib import Path
 from typing import Optional
 
 import click
@@ -10,14 +11,22 @@ from pydantic import BaseModel
 
 from linkml_store import Client
 from linkml_store.api import Collection, Database
+from linkml_store.api.config import ClientConfig
 from linkml_store.api.queries import Query
 from linkml_store.index import get_indexer
 from linkml_store.index.implementations.simple_indexer import SimpleIndexer
 from linkml_store.index.indexer import Indexer
+from linkml_store.inference import get_inference_engine
+from linkml_store.inference.inference_config import InferenceConfig
+from linkml_store.inference.inference_engine import ModelSerialization
 from linkml_store.utils.format_utils import Format, guess_format, load_objects, render_output, write_output
 from linkml_store.utils.object_utils import object_path_update
 from linkml_store.utils.pandas_utils import facet_summary_to_dataframe_unmelted
 
+DEFAULT_LOCAL_CONF_PATH = Path("linkml.yaml")
+# global path is ~/.linkml.yaml in the user's home directory
+DEFAULT_GLOBAL_CONF_PATH = Path("~/.linkml.yaml").expanduser()
+
 index_type_option = click.option(
     "--index-type",
     "-t",
@@ -84,6 +93,7 @@ include_internal_option = click.option("--include-internal/--no-include-internal
 @click.group()
 @click.option("--database", "-d", help="Database name")
 @click.option("--collection", "-c", help="Collection name")
+@click.option("--input", "-i", help="Input file (alternative to database/collection)")
 @click.option("--config", "-C", type=click.Path(exists=True), help="Path to the configuration file")
 @click.option("--set", help="Metadata settings in the form PATHEXPR=value", multiple=True)
 @click.option("-v", "--verbose", count=True)
@@ -96,7 +106,7 @@ include_internal_option = click.option("--include-internal/--no-include-internal
     help="If set then show full stacktrace on error",
 )
 @click.pass_context
-def cli(ctx, verbose: int, quiet: bool, stacktrace: bool, database, collection, config, set, **kwargs):
+def cli(ctx, verbose: int, quiet: bool, stacktrace: bool, database, collection, config, set, input, **kwargs):
     """A CLI for interacting with the linkml-store."""
     if not stacktrace:
         sys.tracebacklimit = 0
@@ -119,13 +129,25 @@ def cli(ctx, verbose: int, quiet: bool, stacktrace: bool, database, collection,
     if quiet:
         logger.setLevel(logging.ERROR)
     ctx.ensure_object(dict)
+    if input:
+        stem = Path(input).stem
+        database = "duckdb"
+        collection = stem
+        config = ClientConfig(databases={"duckdb": {"collections": {stem: {"source": {"local_path": input}}}}})
+        # collection = Path(input).stem
+        # database = f"file:{Path(input).parent}"
+    if config is None and DEFAULT_LOCAL_CONF_PATH.exists():
+        config = DEFAULT_LOCAL_CONF_PATH
+    if config is None and DEFAULT_GLOBAL_CONF_PATH.exists():
+        config = DEFAULT_GLOBAL_CONF_PATH
+    if config == ".":
+        config = None
+    if not collection and database and "::" in database:
+        database, collection = database.split("::")
+
     client = Client().from_config(config, **kwargs) if config else Client()
     settings = ContextSettings(client=client, database_name=database, collection_name=collection)
     ctx.obj["settings"] = settings
-    # DEPRECATED
-    ctx.obj["client"] = client
-    ctx.obj["database"] = database
-    ctx.obj["collection"] = collection
     if settings.database_name:
         db = client.get_database(database)
         if set:
@@ -136,12 +158,6 @@ def cli(ctx, verbose: int, quiet: bool, stacktrace: bool, database, collection,
                 val = yaml.safe_load(val)
                 logger.info(f"Setting {path} to {val}")
                 db.metadata = object_path_update(db.metadata, path, val)
-        # settings.database = db
-        # DEPRECATED
-        ctx.obj["database_obj"] = db
-        if collection:
-            collection_obj = db.get_collection(collection)
-            ctx.obj["collection_obj"] = collection_obj
     if not settings.database_name:
         # if len(client.databases) != 1:
         #    raise ValueError("Database must be specified if there are multiple databases.")
@@ -323,11 +339,12 @@ def apply(ctx, patch_files, identifier_attribute):
 
 @cli.command()
 @click.option("--where", "-w", type=click.STRING, help="WHERE clause for the query, as YAML")
+@click.option("--select", "-s", type=click.STRING, help="SELECT clause for the query, as YAML")
 @click.option("--limit", "-l", type=click.INT, help="Maximum number of results to return")
 @click.option("--output-type", "-O", type=format_choice, default="json", help="Output format")
 @click.option("--output", "-o", type=click.Path(), help="Output file path")
 @click.pass_context
-def query(ctx, where, limit, output_type, output):
+def query(ctx, where, select, limit, output_type, output):
     """Query objects from the specified collection.
 
 
@@ -353,7 +370,13 @@ def query(ctx, where, limit, output_type, output):
     """
     collection = ctx.obj["settings"].collection
     where_clause = yaml.safe_load(where) if where else None
-    query = Query(from_table=collection.alias, where_clause=where_clause, limit=limit)
+    select_clause = yaml.safe_load(select) if select else None
+    if select_clause:
+        if isinstance(select_clause, str):
+            select_clause = [select_clause]
+        if not isinstance(select_clause, list):
+            raise ValueError(f"SELECT clause must be a list. Got: {select_clause}")
+    query = Query(from_table=collection.alias, select_cols=select_clause, where_clause=where_clause, limit=limit)
     result = collection.query(query)
     output_data = render_output(result.rows, output_type)
     if output:
@@ -458,6 +481,110 @@ def describe(ctx, where, output_type, output, limit):
     write_output(df.describe(include="all").transpose(), output_type, target=output)
 
 
+@cli.command()
+@click.option("--output-type", "-O", type=format_choice, default=Format.YAML.value, help="Output format")
+@click.option("--output", "-o", type=click.Path(), help="Output file path")
+@click.option("--target-attribute", "-T", type=click.STRING, multiple=True, help="Target attributes for inference")
+@click.option(
+    "--feature-attributes", "-F", type=click.STRING, help="Feature attributes for inference (comma separated)"
+)
+@click.option("--inference-config-file", "-Y", type=click.Path(), help="Path to inference configuration file")
+@click.option("--export-model", "-E", type=click.Path(), help="Export model to file")
+@click.option("--load-model", "-L", type=click.Path(), help="Load model from file")
+@click.option("--model-format", "-M", type=click.Choice([x.value for x in ModelSerialization]), help="Format for model")
+@click.option("--training-test-data-split", "-S", type=click.Tuple([float, float]), help="Training/test data split")
+@click.option(
+    "--predictor-type", "-t", default="sklearn", show_default=True, type=click.STRING, help="Type of predictor"
+)
+@click.option("--query", "-q", type=click.STRING, help="query term")
+@click.pass_context
+def infer(
+    ctx,
+    inference_config_file,
+    query,
+    training_test_data_split,
+    predictor_type,
+    target_attribute,
+    feature_attributes,
+    output_type,
+    output,
+    model_format,
+    export_model,
+    load_model,
+):
+    """
+    Predict a complete object from a partial object.
+
+    Currently two main prediction methods are provided: RAG and sklearn
+
+    ## RAG:
+
+    The RAG approach will use Retrieval Augmented Generation to inference the missing attributes of an object.
+
+    Example:
+
+        linkml-store  -i countries.jsonl inference -t rag  -q 'name: Uruguay'
+
+    Result:
+
+        capital: Montevideo, code: UY, continent: South America, languages: [Spanish]
+
+    You can pass in configurations as follows:
+
+        linkml-store  -i countries.jsonl inference -t rag:llm_config.model_name=llama-3  -q 'name: Uruguay'
+
+    ## SKLearn:
+
+    This uses scikit-learn (defaulting to simple decision trees) to do the prediction.
+
+        linkml-store -i tests/input/iris.csv inference -t sklearn \
+           -q '{"sepal_length": 5.1, "sepal_width": 3.5, "petal_length": 1.4, "petal_width": 0.2}'
+    """
+    if query:
+        query_obj = yaml.safe_load(query)
+    else:
+        query_obj = None
+    collection = ctx.obj["settings"].collection
+    atts = collection.class_definition().attributes.keys()
+    if model_format:
+        model_format = ModelSerialization(model_format)
+    if load_model:
+        predictor = get_inference_engine(predictor_type)
+        predictor = type(predictor).load_model(load_model)
+    else:
+        if feature_attributes:
+            features = feature_attributes.split(",")
+            features = [f.strip() for f in features]
+        else:
+            if query_obj:
+                features = query_obj.keys()
+            else:
+                features = None
+        if target_attribute:
+            target_attributes = list(target_attribute)
+        else:
+            target_attributes = [att for att in atts if att not in features]
+        if inference_config_file:
+            config = InferenceConfig.from_file(inference_config_file)
+        else:
+            config = InferenceConfig(target_attributes=target_attributes, feature_attributes=features)
+        if training_test_data_split:
+            config.train_test_split = training_test_data_split
+        predictor = get_inference_engine(predictor_type, config=config)
+        predictor.load_and_split_data(collection)
+        predictor.initialize_model()
+    if export_model:
+        logger.info(f"Exporting model to {export_model} in {model_format}")
+        predictor.export_model(export_model, model_format)
+    if not query_obj:
+        if not export_model:
+            raise ValueError("Query must be specified if not exporting model")
+    if query_obj:
+        result = predictor.derive(query_obj)
+        dumped_obj = result.model_dump(exclude_none=True)
+        write_output([dumped_obj], output_type, target=output)
+
+
 @cli.command()
 @index_type_option
 @click.option("--cached-embeddings-database", "-E", help="Path to the database where embeddings are cached")

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/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",
+]

linkml_store/inference/implementations/rag_inference_engine.py

@@ -0,0 +1,145 @@
+import logging
+from dataclasses import dataclass
+from typing import Any, Optional
+
+import yaml
+from llm import get_key
+
+from linkml_store.api.collection import OBJECT, Collection
+from linkml_store.inference.inference_config import Inference, InferenceConfig, LLMConfig
+from linkml_store.inference.inference_engine import InferenceEngine
+
+logger = logging.getLogger(__name__)
+
+SYSTEM_PROMPT = """
+You are a {llm_config.role}, your task is to inference the YAML
+object output given the YAML object input. I will provide you
+with a collection of examples that will provide guidance both
+on the desired structure of the response, as well as the kind
+of content.
+
+You should return ONLY valid YAML in your response.
+"""
+
+
+@dataclass
+class RAGInferenceEngine(InferenceEngine):
+    """
+    AI Retrieval Augmented Generation (RAG) based predictor.
+
+
+    >>> from linkml_store.api.client import Client
+    >>> from linkml_store.utils.format_utils import Format
+    >>> from linkml_store.inference.inference_config import LLMConfig
+    >>> client = Client()
+    >>> db = client.attach_database("duckdb", alias="test")
+    >>> db.import_database("tests/input/countries/countries.jsonl", Format.JSONL, collection_name="countries")
+    >>> db.list_collection_names()
+    ['countries']
+    >>> collection = db.get_collection("countries")
+    >>> features = ["name"]
+    >>> targets = ["code", "capital", "continent", "languages"]
+    >>> llm_config = LLMConfig(model_name="gpt-4o-mini",)
+    >>> config = InferenceConfig(target_attributes=targets, feature_attributes=features, llm_config=llm_config)
+    >>> ie = RAGInferenceEngine(config=config)
+    >>> ie.load_and_split_data(collection)
+    >>> ie.initialize_model()
+    >>> prediction = ie.derive({"name": "Uruguay"})
+    >>> prediction.predicted_object
+    {'capital': 'Montevideo', 'code': 'UY', 'continent': 'South America', 'languages': ['Spanish']}
+
+    """
+
+    classifier: Any = None
+    encoders: dict = None
+    _model: "llm.Model" = None  # noqa: F821
+
+    rag_collection: Collection = None
+
+    def __post_init__(self):
+        if not self.config:
+            self.config = InferenceConfig()
+        if not self.config.llm_config:
+            self.config.llm_config = LLMConfig()
+
+    @property
+    def model(self) -> "llm.Model":  # noqa: F821
+        import llm
+
+        if self._model is None:
+            self._model = llm.get_model(self.config.llm_config.model_name)
+            if self._model.needs_key:
+                key = get_key(None, key_alias=self._model.needs_key)
+                self._model.key = key
+
+        return self._model
+
+    def initialize_model(self, **kwargs):
+        td = self.training_data
+        s = td.slice
+        if not s[0] and not s[1]:
+            rag_collection = td.collection
+        else:
+            base_collection = td.collection
+            objs = base_collection.find({}, offset=s[0], limit=s[1] - s[0]).rows
+            db = base_collection.parent
+            rag_collection = db.get_collection(f"{base_collection.alias}__rag_{s[0]}_{s[1]}", create_if_not_exists=True)
+            rag_collection.insert(objs)
+        rag_collection.attach_indexer("llm", auto_index=False)
+        self.rag_collection = rag_collection
+
+    def object_to_text(self, object: OBJECT) -> str:
+        return yaml.dump(object)
+
+    def derive(self, object: OBJECT) -> Optional[Inference]:
+        import llm
+        from tiktoken import encoding_for_model
+
+        from linkml_store.utils.llm_utils import get_token_limit, render_formatted_text
+
+        model: llm.Model = self.model
+        model_name = self.config.llm_config.model_name
+        feature_attributes = self.config.feature_attributes
+        target_attributes = self.config.target_attributes
+        num_examples = self.config.llm_config.number_of_few_shot_examples or 5
+        query_text = self.object_to_text(object)
+        if not self.rag_collection.indexers:
+            raise ValueError("RAG collection must have an indexer attached")
+        rs = self.rag_collection.search(query_text, limit=num_examples, index_name="llm")
+        examples = rs.rows
+        if not examples:
+            raise ValueError(f"No examples found for {query_text}; size = {self.rag_collection.size()}")
+        prompt_clauses = []
+        for example in examples:
+            input_obj = {k: example.get(k, None) for k in feature_attributes}
+            output_obj = {k: example.get(k, None) for k in target_attributes}
+            prompt_clause = (
+                "---\nExample:\n"
+                f"## INPUT:\n{self.object_to_text(input_obj)}\n"
+                f"## OUTPUT:\n{self.object_to_text(output_obj)}\n"
+            )
+            prompt_clauses.append(prompt_clause)
+        query_obj = {k: object.get(k, None) for k in feature_attributes}
+        query_text = self.object_to_text(query_obj)
+        prompt_end = "---\nQuery:\n" f"## INPUT:\n{query_text}\n" "## OUTPUT:\n"
+        system_prompt = SYSTEM_PROMPT.format(llm_config=self.config.llm_config)
+
+        def make_text(texts):
+            return "\n".join(prompt_clauses) + prompt_end
+
+        try:
+            encoding = encoding_for_model(model_name)
+        except KeyError:
+            encoding = encoding_for_model("gpt-4")
+        token_limit = get_token_limit(model_name)
+        prompt = render_formatted_text(make_text, prompt_clauses, encoding, token_limit)
+        logger.info(f"Prompt: {prompt}")
+        response = model.prompt(prompt, system_prompt)
+        yaml_str = response.text()
+        logger.info(f"Response: {yaml_str}")
+        try:
+            predicted_object = yaml.safe_load(yaml_str)
+            return Inference(predicted_object=predicted_object)
+        except yaml.parser.ParserError as e:
+            logger.error(f"Error parsing response: {yaml_str}\n{e}")
+            return None

linkml_store/inference/implementations/rule_based_inference_engine.py

@@ -0,0 +1,158 @@
+import logging
+from copy import copy
+from dataclasses import dataclass
+from io import StringIO
+from pathlib import Path
+from typing import Any, ClassVar, Dict, List, Optional, Union
+
+import yaml
+from linkml_map.utils.eval_utils import eval_expr
+from linkml_runtime import SchemaView
+from linkml_runtime.linkml_model.meta import AnonymousClassExpression, ClassRule
+from linkml_runtime.utils.formatutils import underscore
+from pydantic import BaseModel
+
+from linkml_store.api.collection import OBJECT, Collection
+from linkml_store.inference.inference_config import Inference
+from linkml_store.inference.inference_engine import InferenceEngine, ModelSerialization
+
+logger = logging.getLogger(__name__)
+
+
+def expression_matches(ce: AnonymousClassExpression, object: OBJECT) -> bool:
+    """
+    Check if a class expression matches an object.
+
+    :param ce: The class expression
+    :param object: The object
+    :return: True if the class expression matches the object
+    """
+    if ce.any_of:
+        if not any(expression_matches(subce, object) for subce in ce.any_of):
+            return False
+    if ce.all_of:
+        if not all(expression_matches(subce, object) for subce in ce.all_of):
+            return False
+    if ce.none_of:
+        if any(expression_matches(subce, object) for subce in ce.none_of):
+            return False
+    if ce.slot_conditions:
+        for slot in ce.slot_conditions.values():
+            slot_name = slot.name
+            v = object.get(slot_name, None)
+            if slot.equals_string is not None:
+                if slot.equals_string != str(v):
+                    return False
+            if slot.equals_integer is not None:
+                if slot.equals_integer != v:
+                    return False
+            if slot.equals_expression is not None:
+                eval_v = eval_expr(slot.equals_expression, **object)
+                if v != eval_v:
+                    return False
+    return True
+
+
+def apply_rule(rule: ClassRule, object: OBJECT):
+    """
+    Apply a rule to an object.
+
+    Mutates the object
+
+    :param rule: The rule to apply
+    :param object: The object to apply the rule to
+    """
+    for condition in rule.preconditions:
+        if expression_matches(condition, object):
+            for postcondition in rule.postconditions:
+                all_of = [x for x in postcondition.all_of] + [postcondition]
+                for pc in all_of:
+                    sc = pc.slot_condition
+                    if sc:
+                        if sc.equals_string:
+                            object[sc.name] = sc.equals_string
+                        if sc.equals_integer:
+                            object[sc.name] = sc.equals_integer
+                        if sc.equals_expression:
+                            object[sc.name] = eval_expr(sc.equals_expression, **object)
+    return object
+
+
+@dataclass
+class RuleBasedInferenceEngine(InferenceEngine):
+    """
+    TODO
+
+    """
+
+    class_rules: Optional[List[ClassRule]] = None
+    slot_rules: Optional[Dict[str, List[ClassRule]]] = None
+    slot_expressions: Optional[Dict[str, str]] = None
+
+    PERSIST_COLS: ClassVar = ["config", "class_rules", "slot_rules", "slot_expressions"]
+
+    def initialize_model(self, **kwargs):
+        td = self.training_data
+        collection: Collection = td.collection
+        cd = collection.class_definition()
+        sv: SchemaView = collection.parent.schema_view
+        class_rules = cd.rules
+        if class_rules:
+            self.class_rules = class_rules
+        for slot in sv.class_induced_slots(cd.name):
+            if slot.equals_expression:
+                self.slot_expressions[slot.name] = slot.equals_expression
+
+    def derive(self, object: OBJECT) -> Optional[Inference]:
+        object = copy(object)
+        if self.class_rules:
+            for rule in self.class_rules:
+                apply_rule(rule, object)
+        object = {underscore(k): v for k, v in object.items()}
+        if self.slot_expressions:
+            for slot, expr in self.slot_expressions.items():
+                print(f"EVAL {object}")
+                v = eval_expr(expr, **object)
+                if v is not None:
+                    object[slot] = v
+        return Inference(predicted_object=object)
+
+    def import_model_from(self, inference_engine: InferenceEngine, **kwargs):
+        io = StringIO()
+        inference_engine.export_model(io, model_serialization=ModelSerialization.LINKML_EXPRESSION)
+        config = inference_engine.config
+        if len(config.target_attributes) != 1:
+            raise ValueError("Can only import models with a single target attribute")
+        target_attribute = config.target_attributes[0]
+        if self.slot_expressions is None:
+            self.slot_expressions = {}
+        self.slot_expressions[target_attribute] = io.getvalue()
+
+    def save_model(self, output: Union[str, Path]) -> None:
+        """
+        Save the trained model and related data to a file.
+
+        :param output: Path to save the model
+        """
+
+        def _serialize_value(v: Any) -> Any:
+            if isinstance(v, BaseModel):
+                return v.model_dump(exclude_unset=True)
+            return v
+
+        model_data = {k: _serialize_value(getattr(self, k)) for k in self.PERSIST_COLS}
+        with open(output, "w", encoding="utf-8") as f:
+            yaml.dump(model_data, f)
+
+    @classmethod
+    def load_model(cls, file_path: Union[str, Path]) -> "RuleBasedInferenceEngine":
+        model_data = yaml.safe_load(open(file_path))
+
+        engine = cls(config=model_data["config"])
+        for k, v in model_data.items():
+            if k == "config":
+                continue
+            setattr(engine, k, v)
+
+        logger.info(f"Model loaded from {file_path}")
+        return engine