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

amsdal_ml/ml_ingesting/splitters/splitter.py

@@ -0,0 +1,15 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+
+from amsdal_ml.ml_ingesting.types import LoadedDocument
+from amsdal_ml.ml_ingesting.types import TextChunk
+
+
+class Splitter(ABC):
+    @abstractmethod
+    def split(self, doc: LoadedDocument) -> list[TextChunk]: ...
+
+    @abstractmethod
+    async def asplit(self, doc: LoadedDocument) -> list[TextChunk]: ...

amsdal_ml/ml_ingesting/splitters/token_splitter.py

@@ -0,0 +1,85 @@
+from __future__ import annotations
+
+import asyncio
+from collections.abc import Callable
+
+from amsdal_ml.ml_ingesting.splitters.splitter import Splitter
+from amsdal_ml.ml_ingesting.types import LoadedDocument
+from amsdal_ml.ml_ingesting.types import TextChunk
+
+
+def _default_token_len(text: str) -> int:
+    return max(1, len(text) // 4)
+
+
+def _words_with_token_counts(text: str, token_fn: Callable[[str], int]) -> list[tuple[str, int]]:
+    out: list[tuple[str, int]] = []
+    for word in text.split():
+        out.append((word, token_fn(word)))
+    return out
+
+
+def _compute_overlap_step(chunk: list[tuple[str, int]], overlap_tokens: int) -> int:
+    tokens = 0
+    step = 0
+    for _word, tk in reversed(chunk):
+        tokens += tk
+        step += 1
+        if tokens >= overlap_tokens:
+            break
+    return step
+
+
+class TokenSplitter(Splitter):
+    def __init__(
+        self,
+        *,
+        max_tokens: int = 600,
+        overlap_tokens: int = 50,
+        token_len_fn: Callable[[str], int] = _default_token_len,
+    ) -> None:
+        self.max_tokens = max_tokens
+        self.overlap_tokens = overlap_tokens
+        self.token_len_fn = token_len_fn
+
+    def _split_page(self, text: str) -> list[str]:
+        words = _words_with_token_counts(text, self.token_len_fn)
+        chunks: list[str] = []
+        start = 0
+        n = len(words)
+        while start < n:
+            tokens = 0
+            end = start
+            while end < n and (tokens + words[end][1] <= self.max_tokens or tokens == 0):
+                tokens += words[end][1]
+                end += 1
+            chunk_words = [w for w, _ in words[start:end]]
+            chunk_text = ' '.join(chunk_words).strip()
+            if chunk_text:
+                chunks.append(chunk_text)
+            if end >= n:
+                break
+            back = _compute_overlap_step(words[start:end], self.overlap_tokens)
+            start = max(start + 1, end - back)
+        return chunks
+
+    def split(self, doc: LoadedDocument) -> list[TextChunk]:
+        chunks: list[TextChunk] = []
+        for page in doc.pages:
+            page_chunks = self._split_page(page.text)
+            for chunk_text in page_chunks:
+                metadata = dict(doc.metadata)
+                metadata.update(page.metadata)
+                metadata.setdefault('page_number', page.page_number)
+                chunks.append(
+                    TextChunk(
+                        index=len(chunks),
+                        text=chunk_text,
+                        metadata=metadata,
+                        tags=[],
+                    )
+                )
+        return chunks
+
+    async def asplit(self, doc: LoadedDocument) -> list[TextChunk]:
+        return await asyncio.to_thread(self.split, doc)

amsdal_ml/ml_ingesting/stores/__init__.py

@@ -0,0 +1,4 @@
+from amsdal_ml.ml_ingesting.stores.embedding_data import EmbeddingDataStore
+from amsdal_ml.ml_ingesting.stores.store import EmbeddingStore
+
+__all__ = ['EmbeddingDataStore', 'EmbeddingStore']

amsdal_ml/ml_ingesting/stores/embedding_data.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+from typing import cast
+
+from amsdal_ml.ml_ingesting.embedding_data import EmbeddingData
+from amsdal_ml.ml_ingesting.stores.store import EmbeddingStore
+from amsdal_ml.ml_ingesting.types import IngestionSource
+from amsdal_ml.models.embedding_model import EmbeddingModel
+
+
+class EmbeddingDataStore(EmbeddingStore):
+    def __init__(self, *, model_cls=EmbeddingModel) -> None:
+        self.model_cls = model_cls
+
+    def _merge_tags(self, base: list[str], extra: list[str]) -> list[str]:
+        merged = list(base)
+        for tag in extra:
+            if tag not in merged:
+                merged.append(tag)
+        return merged
+
+    def _build_objects(self, embeddings: Sequence[EmbeddingData], source: IngestionSource) -> list[Any]:
+        base_tags = list(source.tags)
+        base_meta = dict(source.metadata)
+        objs = []
+        for record in embeddings:
+            tags = self._merge_tags(base_tags, list(record.tags))
+            meta = dict(base_meta)
+            meta.update(record.metadata or {})
+            objs.append(
+                self.model_cls(
+                    data_object_class=source.object_class,
+                    data_object_id=source.object_id,
+                    chunk_index=record.chunk_index,
+                    raw_text=record.raw_text,
+                    embedding=record.embedding,
+                    tags=tags,
+                    ml_metadata=meta,
+                )
+            )
+        return objs
+
+    def save(self, embeddings: Sequence[EmbeddingData], *, source: IngestionSource | None = None) -> list[Any]:
+        if source is None:
+            msg = 'source is required to save embeddings'
+            raise RuntimeError(msg)
+        objs = self._build_objects(embeddings, source)
+        manager = cast(Any, self.model_cls.objects)
+        manager.bulk_create(objs)
+        return objs
+
+    async def asave(
+        self, embeddings: Sequence[EmbeddingData], *, source: IngestionSource | None = None
+    ) -> list[Any]:
+        if source is None:
+            msg = 'source is required to save embeddings'
+            raise RuntimeError(msg)
+        objs = self._build_objects(embeddings, source)
+        manager = cast(Any, self.model_cls.objects)
+        await manager.bulk_acreate(objs)
+        return objs

amsdal_ml/ml_ingesting/stores/store.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Sequence
+from typing import Any
+
+from amsdal_ml.ml_ingesting.embedding_data import EmbeddingData
+from amsdal_ml.ml_ingesting.types import IngestionSource
+
+
+class EmbeddingStore(ABC):
+    @abstractmethod
+    def save(self, embeddings: Sequence[EmbeddingData], *, source: IngestionSource | None = None) -> list[Any]: ...
+
+    @abstractmethod
+    async def asave(
+        self,
+        embeddings: Sequence[EmbeddingData],
+        *,
+        source: IngestionSource | None = None,
+    ) -> list[Any]: ...

amsdal_ml/ml_ingesting/types.py

@@ -0,0 +1,40 @@
+from __future__ import annotations
+
+from typing import Any
+
+from pydantic import BaseModel
+from pydantic import Field
+
+
+class LoadedPage(BaseModel):
+    page_number: int | None = Field(default=None, title='1-based page number if available')
+    text: str = Field(..., min_length=1, title='Extracted page text')
+    metadata: dict[str, Any] = Field(default_factory=dict, title='Page-level metadata')
+
+    def as_text(self) -> str:
+        return self.text
+
+
+class LoadedDocument(BaseModel):
+    pages: list[LoadedPage] = Field(default_factory=list, title='Pages in original order')
+    metadata: dict[str, Any] = Field(default_factory=dict, title='Document-level metadata')
+
+    def join(self, *, separator: str = '\n\n') -> str:
+        return separator.join(page.text for page in self.pages)
+
+
+class TextChunk(BaseModel):
+    index: int = Field(..., title='Chunk order in document')
+    text: str = Field(..., min_length=1, title='Chunk text destined for embedding')
+    tags: list[str] = Field(default_factory=list, title='Tags to persist with embeddings')
+    metadata: dict[str, Any] = Field(default_factory=dict, title='Arbitrary chunk metadata')
+
+
+class IngestionSource(BaseModel):
+    object_class: str = Field(..., title='Linked object class for embeddings')
+    object_id: str = Field(..., title='Linked object ID for embeddings')
+    tags: list[str] = Field(default_factory=list, title='Base tags applied to all chunks')
+    metadata: dict[str, Any] = Field(default_factory=dict, title='Arbitrary source metadata')
+
+
+__all__ = ['IngestionSource', 'LoadedDocument', 'LoadedPage', 'TextChunk']

amsdal_ml/ml_models/models.py

@@ -3,9 +3,31 @@ from __future__ import annotations
 from abc import ABC
 from abc import abstractmethod
 from collections.abc import AsyncIterator
+from typing import Any
+from typing import Required
+from typing import TypedDict
 
 from amsdal_ml.fileio.base_loader import PLAIN_TEXT
 from amsdal_ml.fileio.base_loader import FileAttachment
+from amsdal_ml.ml_models.utils import ResponseFormat
+
+
+class StructuredMessage(TypedDict, total=False):
+    """Base structure for a message in LLM conversations.
+
+    Attributes:
+        role: The role of the message sender (e.g., 'user', 'assistant', 'system').
+        content: The content of the message, can be str or list of multimodal parts.
+        tool_call_id: ID of the tool call (for tool messages).
+        name: Name of the tool (for tool messages).
+    """
+    role: Required[str]
+    content: Required[Any]
+    tool_call_id: str
+    name: str
+
+
+LLModelInput = str | list[StructuredMessage]
 
 
 class ModelError(Exception):
@@ -25,6 +47,12 @@ class ModelAPIError(ModelError):
 
 
 class MLModel(ABC):
+    @property
+    @abstractmethod
+    def supported_formats(self) -> set[ResponseFormat]:
+        """Return a set of supported response formats for this model."""
+        raise NotImplementedError
+
     @abstractmethod
     def setup(self) -> None:
         """Initialize any clients or resources needed before inference."""
@@ -39,12 +67,64 @@ class MLModel(ABC):
         """Return a set of universal attachment kinds, e.g. {PLAIN_TEXT, FILE_ID}."""
         return {PLAIN_TEXT}
 
+    @property
+    @abstractmethod
+    def input_role(self) -> str:
+        """Return the role for user input messages."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def output_role(self) -> str:
+        """Return the role for model output messages."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def tool_role(self) -> str:
+        """Return the role for tool result messages."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def system_role(self) -> str:
+        """Return the role for system messages."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def content_field(self) -> str:
+        """Return the field name for message content (e.g., 'content' for OpenAI, 'parts' for Gemini)."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def role_field(self) -> str:
+        """Return the field name for message role (e.g., 'role' for most models)."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def tool_call_id_field(self) -> str:
+        """Return the field name for tool call ID (e.g., 'tool_call_id' for OpenAI)."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def tool_name_field(self) -> str:
+        """Return the field name for tool name (e.g., 'name' for OpenAI)."""
+        raise NotImplementedError
+
     @abstractmethod
     def invoke(
         self,
-        prompt: str,
+        input: LLModelInput,  # noqa: A002
         *,
         attachments: list[FileAttachment] | None = None,
+        response_format: ResponseFormat | None = None,
+        schema: dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str | dict[str, Any] | None = None,
     ) -> str:
         """Run synchronous inference with the model."""
         raise NotImplementedError
@@ -52,9 +132,13 @@ class MLModel(ABC):
     @abstractmethod
     async def ainvoke(
         self,
-        prompt: str,
+        input: LLModelInput,  # noqa: A002
         *,
         attachments: list[FileAttachment] | None = None,
+        response_format: ResponseFormat | None = None,
+        schema: dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str | dict[str, Any] | None = None,
     ) -> str:
         """Run asynchronous inference with the model."""
         raise NotImplementedError
@@ -62,9 +146,13 @@ class MLModel(ABC):
     @abstractmethod
     def stream(
         self,
-        prompt: str,
+        input: LLModelInput,  # noqa: A002
         *,
         attachments: list[FileAttachment] | None = None,
+        response_format: ResponseFormat | None = None,
+        schema: dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str | dict[str, Any] | None = None,
     ):
         """Stream synchronous inference results from the model."""
         raise NotImplementedError
@@ -72,9 +160,13 @@ class MLModel(ABC):
     @abstractmethod
     def astream(
         self,
-        prompt: str,
+        input: LLModelInput,  # noqa: A002
         *,
         attachments: list[FileAttachment] | None = None,
+        response_format: ResponseFormat | None = None,
+        schema: dict[str, Any] | None = None,
+        tools: list[dict[str, Any]] | None = None,
+        tool_choice: str | dict[str, Any] | None = None,
     ) -> AsyncIterator[str]:
         """
         Stream asynchronous inference results as an async generator.