draive 0.9.1__py3-none-any.whl

draive/__init__.py

@@ -0,0 +1,186 @@
+from draive.conversation import (
+    Conversation,
+    ConversationCompletion,
+    ConversationCompletionStream,
+    ConversationMessage,
+    ConversationMessageContent,
+    conversation_completion,
+    lmm_conversation_completion,
+)
+from draive.embedding import Embedded, Embedder, Embedding, embed_text
+from draive.generation import (
+    ImageGeneration,
+    ImageGenerator,
+    ModelGeneration,
+    ModelGenerator,
+    TextGeneration,
+    TextGenerator,
+    generate_image,
+    generate_model,
+    generate_text,
+)
+from draive.helpers import (
+    getenv_bool,
+    getenv_float,
+    getenv_int,
+    getenv_str,
+    load_env,
+    setup_logging,
+    split_sequence,
+)
+from draive.lmm import (
+    LMM,
+    LMMCompletion,
+    LMMCompletionContent,
+    LMMCompletionMessage,
+    LMMCompletionStream,
+    LMMCompletionStreamingUpdate,
+    lmm_completion,
+)
+from draive.mistral import (
+    MistralChatConfig,
+    MistralClient,
+    MistralException,
+    mistral_lmm_completion,
+)
+from draive.openai import (
+    OpenAIChatConfig,
+    OpenAIClient,
+    OpenAIEmbeddingConfig,
+    OpenAIException,
+    OpenAIImageGenerationConfig,
+    openai_embed_text,
+    openai_generate_image,
+    openai_lmm_completion,
+    openai_tokenize_text,
+)
+from draive.scope import (
+    DependenciesScope,
+    MetricsScope,
+    ScopeDependency,
+    ScopeMetric,
+    StateScope,
+    TokenUsage,
+    ctx,
+)
+from draive.similarity import mmr_similarity, similarity
+from draive.splitters import split_text
+from draive.tokenization import TextTokenizer, Tokenization, count_text_tokens, tokenize_text
+from draive.tools import (
+    Tool,
+    Toolbox,
+    ToolCallContext,
+    ToolCallStatus,
+    ToolCallUpdate,
+    ToolException,
+    ToolsUpdatesContext,
+    tool,
+)
+from draive.types import (
+    MISSING,
+    Argument,
+    Field,
+    ImageBase64Content,
+    ImageContent,
+    ImageURLContent,
+    Memory,
+    MissingValue,
+    Model,
+    MultimodalContent,
+    ReadOnlyMemory,
+    State,
+    UpdateSend,
+)
+from draive.utils import allowing_early_exit, auto_retry, cache, with_early_exit
+
+__all__ = [
+    "Conversation",
+    "ImageContent",
+    "ImageBase64Content",
+    "ImageURLContent",
+    "UpdateSend",
+    "Embedded",
+    "Embedder",
+    "Embedding",
+    "Model",
+    "Memory",
+    "ModelGeneration",
+    "ModelGenerator",
+    "ReadOnlyMemory",
+    "DependenciesScope",
+    "ScopeDependency",
+    "ScopeMetric",
+    "MetricsScope",
+    "StateScope",
+    "State",
+    "Field",
+    "TextGeneration",
+    "TextGenerator",
+    "ImageGeneration",
+    "ImageGenerator",
+    "generate_image",
+    "TextTokenizer",
+    "Tokenization",
+    "TokenUsage",
+    "auto_retry",
+    "cache",
+    "conversation_completion",
+    "tokenize_text",
+    "count_text_tokens",
+    "ctx",
+    "embed_text",
+    "generate_model",
+    "generate_text",
+    "load_env",
+    "getenv_bool",
+    "getenv_float",
+    "getenv_int",
+    "getenv_str",
+    "mmr_similarity",
+    "similarity",
+    "split_sequence",
+    "split_text",
+    "tool",
+    "Argument",
+    "allowing_early_exit",
+    "with_early_exit",
+    "setup_logging",
+    "MissingValue",
+    "MISSING",
+    "MultimodalContent",
+    "Tool",
+    "ToolException",
+    "Toolbox",
+    "ToolCallContext",
+    "ToolCallStatus",
+    "ToolCallUpdate",
+    "ToolException",
+    "ToolsUpdatesContext",
+    "LMM",
+    "LMMCompletion",
+    "LMMCompletionContent",
+    "LMMCompletionMessage",
+    "LMMCompletionStreamingUpdate",
+    "LMMCompletionStream",
+    "lmm_completion",
+    "ConversationMessage",
+    "ConversationMessageContent",
+    "Conversation",
+    "ConversationCompletionStream",
+    "ConversationCompletion",
+    "conversation_completion",
+    "lmm_conversation_completion",
+    "OpenAIException",
+    "OpenAIChatConfig",
+    "OpenAIClient",
+    "OpenAIEmbeddingConfig",
+    "OpenAIImageGenerationConfig",
+    "openai_tokenize_text",
+    "openai_embed_text",
+    "openai_lmm_completion",
+    "openai_generate_image",
+    "MistralException",
+    "MistralClient",
+    "MistralChatConfig",
+    "mistral_lmm_completion",
+]

draive/conversation/__init__.py

@@ -0,0 +1,18 @@
+from draive.conversation.call import conversation_completion
+from draive.conversation.completion import ConversationCompletion, ConversationCompletionStream
+from draive.conversation.lmm import lmm_conversation_completion
+from draive.conversation.message import (
+    ConversationMessage,
+    ConversationMessageContent,
+)
+from draive.conversation.state import Conversation
+
+__all__ = [
+    "ConversationMessage",
+    "ConversationMessageContent",
+    "Conversation",
+    "ConversationCompletionStream",
+    "ConversationCompletion",
+    "conversation_completion",
+    "lmm_conversation_completion",
+]

draive/conversation/call.py

@@ -0,0 +1,87 @@
+from typing import Literal, overload
+
+from draive.conversation.completion import ConversationCompletionStream
+from draive.conversation.message import (
+    ConversationMessage,
+    ConversationMessageContent,
+    ConversationStreamingUpdate,
+)
+from draive.conversation.state import Conversation
+from draive.scope import ctx
+from draive.tools import Toolbox
+from draive.types import Memory, UpdateSend
+
+__all__ = [
+    "conversation_completion",
+]
+
+
+@overload
+async def conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+    stream: Literal[True],
+) -> ConversationCompletionStream:
+    ...
+
+
+@overload
+async def conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+    stream: UpdateSend[ConversationStreamingUpdate],
+) -> ConversationMessage:
+    ...
+
+
+@overload
+async def conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+) -> ConversationMessage:
+    ...
+
+
+async def conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+    stream: UpdateSend[ConversationStreamingUpdate] | bool = False,
+) -> ConversationCompletionStream | ConversationMessage:
+    conversation: Conversation = ctx.state(Conversation)
+
+    match stream:
+        case False:
+            return await conversation.completion(
+                instruction=instruction,
+                input=input,
+                memory=memory or conversation.memory,
+                tools=tools or conversation.tools,
+            )
+        case True:
+            return await conversation.completion(
+                instruction=instruction,
+                input=input,
+                memory=memory or conversation.memory,
+                tools=tools or conversation.tools,
+                stream=True,
+            )
+        case progress:
+            return await conversation.completion(
+                instruction=instruction,
+                input=input,
+                memory=memory or conversation.memory,
+                tools=tools or conversation.tools,
+                stream=progress,
+            )

draive/conversation/completion.py

@@ -0,0 +1,67 @@
+from typing import Literal, Protocol, overload, runtime_checkable
+
+from draive.conversation.message import (
+    ConversationMessage,
+    ConversationMessageContent,
+    ConversationStreamingUpdate,
+)
+from draive.lmm import LMMCompletionStream
+from draive.tools import Toolbox
+from draive.types import Memory, UpdateSend
+
+__all__ = [
+    "ConversationCompletionStream",
+    "ConversationCompletion",
+]
+
+
+ConversationCompletionStream = LMMCompletionStream
+
+
+@runtime_checkable
+class ConversationCompletion(Protocol):
+    @overload
+    async def __call__(
+        self,
+        *,
+        instruction: str,
+        input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+        memory: Memory[ConversationMessage] | None = None,
+        tools: Toolbox | None = None,
+        stream: Literal[True],
+    ) -> ConversationCompletionStream:
+        ...
+
+    @overload
+    async def __call__(
+        self,
+        *,
+        instruction: str,
+        input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+        memory: Memory[ConversationMessage] | None = None,
+        tools: Toolbox | None = None,
+        stream: UpdateSend[ConversationStreamingUpdate],
+    ) -> ConversationMessage:
+        ...
+
+    @overload
+    async def __call__(
+        self,
+        *,
+        instruction: str,
+        input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+        memory: Memory[ConversationMessage] | None = None,
+        tools: Toolbox | None = None,
+    ) -> ConversationMessage:
+        ...
+
+    async def __call__(  # noqa: PLR0913
+        self,
+        *,
+        instruction: str,
+        input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+        memory: Memory[ConversationMessage] | None = None,
+        tools: Toolbox | None = None,
+        stream: UpdateSend[ConversationStreamingUpdate] | bool = False,
+    ) -> ConversationCompletionStream | ConversationMessage:
+        ...

draive/conversation/lmm.py

@@ -0,0 +1,159 @@
+from datetime import UTC, datetime
+from typing import Literal, overload
+
+from draive.conversation.completion import ConversationCompletionStream
+from draive.conversation.message import (
+    ConversationMessage,
+    ConversationMessageContent,
+    ConversationStreamingUpdate,
+)
+from draive.lmm import LMMCompletionMessage, lmm_completion
+from draive.tools import Toolbox
+from draive.types import Memory, UpdateSend
+from draive.utils import AsyncStreamTask
+
+__all__: list[str] = [
+    "lmm_conversation_completion",
+]
+
+
+@overload
+async def lmm_conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+    stream: Literal[True],
+) -> ConversationCompletionStream:
+    ...
+
+
+@overload
+async def lmm_conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+    stream: UpdateSend[ConversationStreamingUpdate],
+) -> ConversationMessage:
+    ...
+
+
+@overload
+async def lmm_conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+) -> ConversationMessage:
+    ...
+
+
+async def lmm_conversation_completion(
+    *,
+    instruction: str,
+    input: ConversationMessage | ConversationMessageContent,  # noqa: A002
+    memory: Memory[ConversationMessage] | None = None,
+    tools: Toolbox | None = None,
+    stream: UpdateSend[ConversationStreamingUpdate] | bool = False,
+) -> ConversationCompletionStream | ConversationMessage:
+    system_message: LMMCompletionMessage = LMMCompletionMessage(
+        role="system",
+        content=instruction,
+    )
+    user_message: ConversationMessage
+    if isinstance(input, ConversationMessage):
+        user_message = input
+
+    else:
+        user_message = ConversationMessage(
+            timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S%z"),
+            role="user",
+            content=input,
+        )
+
+    context: list[LMMCompletionMessage]
+
+    if memory:
+        context = [
+            system_message,
+            *await memory.recall(),
+            user_message,
+        ]
+
+    else:
+        context = [
+            system_message,
+            user_message,
+        ]
+
+    match stream:
+        case True:
+
+            async def stream_task(
+                update: UpdateSend[ConversationStreamingUpdate],
+            ) -> None:
+                nonlocal memory
+                completion: LMMCompletionMessage = await lmm_completion(
+                    context=context,
+                    tools=tools,
+                    stream=update,
+                )
+                response_message: ConversationMessage = ConversationMessage(
+                    timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S%z"),
+                    role=completion.role,
+                    content=completion.content,
+                )
+                if memory := memory:
+                    await memory.remember(
+                        [
+                            user_message,
+                            response_message,
+                        ],
+                    )
+
+            return AsyncStreamTask(job=stream_task)
+
+        case False:
+            completion: LMMCompletionMessage = await lmm_completion(
+                context=context,
+                tools=tools,
+            )
+            response_message: ConversationMessage = ConversationMessage(
+                timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S%z"),
+                role=completion.role,
+                content=completion.content,
+            )
+            if memory := memory:
+                await memory.remember(
+                    [
+                        user_message,
+                        response_message,
+                    ],
+                )
+
+            return response_message
+
+        case update:
+            completion: LMMCompletionMessage = await lmm_completion(
+                context=context,
+                tools=tools,
+                stream=update,
+            )
+            response_message: ConversationMessage = ConversationMessage(
+                timestamp=datetime.now(UTC).strftime("%Y-%m-%dT%H:%M:%S%z"),
+                role=completion.role,
+                content=completion.content,
+            )
+            if memory := memory:
+                await memory.remember(
+                    [
+                        user_message,
+                        response_message,
+                    ],
+                )
+
+            return response_message

draive/conversation/message.py

@@ -0,0 +1,16 @@
+from draive.lmm import LMMCompletionContent, LMMCompletionMessage, LMMCompletionStreamingUpdate
+
+__all__ = [
+    "ConversationMessage",
+    "ConversationMessageContent",
+    "ConversationStreamingUpdate",
+]
+
+
+class ConversationMessage(LMMCompletionMessage):
+    author: str | None = None
+    timestamp: str | None = None
+
+
+ConversationMessageContent = LMMCompletionContent
+ConversationStreamingUpdate = LMMCompletionStreamingUpdate

draive/conversation/state.py

@@ -0,0 +1,17 @@
+from draive.conversation.completion import ConversationCompletion
+from draive.conversation.lmm import lmm_conversation_completion
+from draive.conversation.message import (
+    ConversationMessage,
+)
+from draive.tools import Toolbox
+from draive.types import Memory, State
+
+__all__: list[str] = [
+    "Conversation",
+]
+
+
+class Conversation(State):
+    completion: ConversationCompletion = lmm_conversation_completion
+    memory: Memory[ConversationMessage] | None = None
+    tools: Toolbox | None = None

draive/embedding/__init__.py

@@ -0,0 +1,11 @@
+from draive.embedding.call import embed_text
+from draive.embedding.embedded import Embedded
+from draive.embedding.embedder import Embedder
+from draive.embedding.state import Embedding
+
+__all__ = [
+    "Embedder",
+    "Embedding",
+    "Embedded",
+    "embed_text",
+]

draive/embedding/call.py

@@ -0,0 +1,15 @@
+from collections.abc import Iterable
+
+from draive.embedding.embedded import Embedded
+from draive.embedding.state import Embedding
+from draive.scope import ctx
+
+__all__ = [
+    "embed_text",
+]
+
+
+async def embed_text(
+    values: Iterable[str],
+) -> list[Embedded[str]]:
+    return await ctx.state(Embedding).embed_text(values=values)

draive/embedding/embedded.py

@@ -0,0 +1,18 @@
+from typing import Generic, TypeVar
+
+from draive.types.state import State
+
+__all__ = [
+    "Embedded",
+]
+
+
+_Embedded = TypeVar(
+    "_Embedded",
+    bound=object,
+)
+
+
+class Embedded(State, Generic[_Embedded]):
+    value: _Embedded
+    vector: list[float]

draive/embedding/embedder.py

@@ -0,0 +1,23 @@
+from collections.abc import Iterable
+from typing import Generic, Protocol, TypeVar, runtime_checkable
+
+from draive.embedding.embedded import Embedded
+
+__all__ = [
+    "Embedder",
+]
+
+
+_Embeddable = TypeVar(
+    "_Embeddable",
+    bound=object,
+)
+
+
+@runtime_checkable
+class Embedder(Protocol, Generic[_Embeddable]):
+    async def __call__(
+        self,
+        values: Iterable[_Embeddable],
+    ) -> list[Embedded[_Embeddable]]:
+        ...

draive/embedding/state.py

@@ -0,0 +1,10 @@
+from draive.embedding.embedder import Embedder
+from draive.types import State
+
+__all__ = [
+    "Embedding",
+]
+
+
+class Embedding(State):
+    embed_text: Embedder[str]

draive/generation/__init__.py

@@ -0,0 +1,15 @@
+from draive.generation.image import ImageGeneration, ImageGenerator, generate_image
+from draive.generation.model import ModelGeneration, ModelGenerator, generate_model
+from draive.generation.text import TextGeneration, TextGenerator, generate_text
+
+__all__ = [
+    "ModelGenerator",
+    "ModelGeneration",
+    "generate_model",
+    "TextGenerator",
+    "TextGeneration",
+    "generate_text",
+    "ImageGeneration",
+    "ImageGenerator",
+    "generate_image",
+]

draive/generation/image/__init__.py

@@ -0,0 +1,9 @@
+from draive.generation.image.call import generate_image
+from draive.generation.image.generator import ImageGenerator
+from draive.generation.image.state import ImageGeneration
+
+__all__ = [
+    "generate_image",
+    "ImageGenerator",
+    "ImageGeneration",
+]

draive/generation/image/call.py

@@ -0,0 +1,16 @@
+from draive.generation.image.state import ImageGeneration
+from draive.scope import ctx
+from draive.types import ImageContent
+
+__all__ = [
+    "generate_image",
+]
+
+
+async def generate_image(
+    *,
+    instruction: str,
+) -> ImageContent:
+    return await ctx.state(ImageGeneration).generate(
+        instruction=instruction,
+    )

draive/generation/image/generator.py

@@ -0,0 +1,17 @@
+from typing import Protocol, runtime_checkable
+
+from draive.types import ImageContent
+
+__all__ = [
+    "ImageGenerator",
+]
+
+
+@runtime_checkable
+class ImageGenerator(Protocol):
+    async def __call__(
+        self,
+        *,
+        instruction: str,
+    ) -> ImageContent:
+        ...