amsdal_ml 0.1.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

amsdal_ml/ml_models/utils.py

@@ -0,0 +1,7 @@
+from enum import Enum
+
+
+class ResponseFormat(Enum):
+    PLAIN_TEXT = "plain_text"
+    JSON_OBJECT = "json_object"
+    JSON_SCHEMA = "json_schema"

amsdal_ml/ml_retrievers/__init__.py

@@ -0,0 +1,17 @@
+from amsdal_ml.ml_retrievers.default_retriever import DefaultRetriever
+from amsdal_ml.ml_retrievers.openai_retriever import OpenAIRetriever
+from amsdal_ml.ml_retrievers.query_retriever import NLQueryExecutor
+from amsdal_ml.ml_retrievers.query_retriever import NLQueryRetriever
+from amsdal_ml.ml_retrievers.retriever import MLRetriever
+from amsdal_ml.ml_retrievers.retriever import RetrievalChunk
+from amsdal_ml.ml_retrievers.retriever import Retriever
+
+__all__ = [
+    "DefaultRetriever",
+    "MLRetriever",
+    "NLQueryExecutor",
+    "NLQueryRetriever",
+    "OpenAIRetriever",
+    "RetrievalChunk",
+    "Retriever",
+]

amsdal_ml/ml_retrievers/adapters.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import json
+from abc import ABC
+from abc import abstractmethod
+from typing import TYPE_CHECKING
+from typing import Any
+
+if TYPE_CHECKING:
+    from amsdal_ml.ml_models.models import MLModel
+    from amsdal_ml.ml_retrievers.query_retriever import FilterCondition
+
+
+class RetrieverAdapter(ABC):
+
+    @abstractmethod
+    def get_response_schema(self, base_schema: dict[str, Any]) -> dict[str, Any]:
+        """Adapts the base JSON schema for the specific model."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def parse_response(
+        self,
+        raw_json: str,
+        *,
+        is_schema_based: bool,
+    ) -> list[FilterCondition]:
+        """Parses the raw JSON string from the model into a list of FilterCondition."""
+        raise NotImplementedError
+
+
+class DefaultRetrieverAdapter(RetrieverAdapter):
+    def get_response_schema(self, base_schema: dict[str, Any]) -> dict[str, Any]:
+        return base_schema
+
+    def parse_response(
+        self,
+        raw_json: str,
+        *,
+        is_schema_based: bool,
+    ) -> list[FilterCondition]:
+        from amsdal_ml.ml_retrievers.query_retriever import FilterCondition
+        from amsdal_ml.ml_retrievers.query_retriever import FilterResponse
+
+        try:
+            filter_data = json.loads(raw_json)
+
+            if is_schema_based:
+                return FilterResponse.model_validate(filter_data).filters
+
+            if isinstance(filter_data, dict) and "filters" in filter_data:
+                return [FilterCondition(**cond) for cond in filter_data["filters"]]
+            if isinstance(filter_data, list):
+                return [FilterCondition(**cond) for cond in filter_data]
+
+            return []
+        except (json.JSONDecodeError, Exception):
+            return []
+
+
+class OpenAIRetrieverAdapter(DefaultRetrieverAdapter):
+    def get_response_schema(self, base_schema: dict[str, Any]) -> dict[str, Any]:
+        def add_additional_properties_recursively(schema_node: dict[str, Any] | list[Any]) -> None:
+            if isinstance(schema_node, dict):
+                if (
+                    schema_node.get("type") == "object"
+                    and "additionalProperties" not in schema_node
+                ):
+                    schema_node["additionalProperties"] = False
+
+                for value in schema_node.values():
+                    add_additional_properties_recursively(value)
+
+            elif isinstance(schema_node, list):
+                for item in schema_node:
+                    add_additional_properties_recursively(item)
+
+        add_additional_properties_recursively(base_schema)
+
+        return {
+            "name": "data",
+            "strict": True,
+            "schema": base_schema
+        }
+
+
+def get_retriever_adapter(model: MLModel) -> RetrieverAdapter:
+    model_name = model.__class__.__name__.lower()
+
+    if "openai" in model_name:
+        return OpenAIRetrieverAdapter()
+
+    return DefaultRetrieverAdapter()

amsdal_ml/ml_retrievers/default_retriever.py

@@ -6,6 +6,7 @@ from abc import abstractmethod
 
 from amsdal_models.classes.annotations import CosineDistance
 
+from amsdal_ml.ml_ingesting.embedders.embedder import Embedder
 from amsdal_ml.models.embedding_model import EmbeddingModel
 
 from .retriever import MLRetriever
@@ -18,14 +19,21 @@ def _default_num_tokens(text: str) -> int:
 
 class DefaultRetriever(MLRetriever, ABC):
     def __init__(
-        self, *, embedding_model_cls=EmbeddingModel, max_context_tokens: int = 1800, num_tokens_fn=_default_num_tokens
+        self,
+        *,
+        embedding_model_cls=EmbeddingModel,
+        max_context_tokens: int = 1800,
+        num_tokens_fn=_default_num_tokens,
+        embedder: Embedder | None = None,
     ):
         self.embedding_model_cls = embedding_model_cls
         self.max_context_tokens = max_context_tokens
         self.num_tokens_fn = num_tokens_fn
+        self.embedder = embedder
 
     @abstractmethod
     def _embed_query(self, text: str) -> list[float]: ...
+
     @abstractmethod
     async def _aembed_query(self, text: str) -> list[float]: ...
 
@@ -66,6 +74,7 @@ class DefaultRetriever(MLRetriever, ABC):
                 raw_text=(r.raw_text or '').strip(),
                 distance=float(getattr(r, 'distance', math.inf)),
                 tags=list(r.tags or []),
+                metadata=dict(getattr(r, 'ml_metadata', {}) or {}),
             )
             for r in rows
         ]
@@ -100,6 +109,7 @@ class DefaultRetriever(MLRetriever, ABC):
                 raw_text=(r.raw_text or '').strip(),
                 distance=float(getattr(r, 'distance', math.inf)),
                 tags=list(r.tags or []),
+                metadata=dict(getattr(r, 'ml_metadata', {}) or {}),
             )
             for r in rows
         ]

amsdal_ml/ml_retrievers/openai_retriever.py

@@ -7,6 +7,7 @@ from openai import AsyncOpenAI
 from openai import OpenAI
 
 from amsdal_ml.ml_config import ml_config
+from amsdal_ml.ml_ingesting.embedders.embedder import Embedder
 
 from .default_retriever import DefaultRetriever
 
@@ -14,26 +15,45 @@ DEFAULT_EMBED_MODEL = ml_config.embed_model_name
 
 
 class OpenAIRetriever(DefaultRetriever):
+    client: OpenAI | None = None
+    aclient: AsyncOpenAI | None = None
+
     def __init__(
         self,
         *,
         api_key: Optional[str] = None,
         embed_model: Optional[str] = None,
+        embedder: Embedder | None = None,
     ):
-        super().__init__()
-        self.api_key = api_key or ml_config.resolved_openai_key or os.getenv('OPENAI_API_KEY')
-        if not self.api_key:
-            msg = 'OPENAI_API_KEY is required for OpenAIRetriever'
-            raise RuntimeError(msg)
+        super().__init__(embedder=embedder)
 
+        self.api_key = api_key or ml_config.resolved_openai_key or os.getenv('OPENAI_API_KEY')
         self.embed_model = embed_model or DEFAULT_EMBED_MODEL
-        self.client = OpenAI(api_key=self.api_key)
-        self.aclient = AsyncOpenAI(api_key=self.api_key)
+
+        if embedder is None:
+            if not self.api_key:
+                msg = 'OPENAI_API_KEY is required for OpenAIRetriever'
+                raise RuntimeError(msg)
+            self.client = OpenAI(api_key=self.api_key)
+            self.aclient = AsyncOpenAI(api_key=self.api_key)
+        else:
+            self.client = None
+            self.aclient = None
 
     def _embed_query(self, text: str) -> list[float]:
+        if self.embedder is not None:
+            return self.embedder.embed(text)
+        if not self.client:
+            msg = 'OpenAI client is not configured'
+            raise RuntimeError(msg)
         resp = self.client.embeddings.create(model=self.embed_model, input=text)
         return resp.data[0].embedding
 
     async def _aembed_query(self, text: str) -> list[float]:
+        if self.embedder is not None:
+            return await self.embedder.aembed(text)
+        if not self.aclient:
+            msg = 'Async OpenAI client is not configured'
+            raise RuntimeError(msg)
         resp = await self.aclient.embeddings.create(model=self.embed_model, input=text)
         return resp.data[0].embedding