seekrai 0.4.4__py3-none-any.whl → 0.5.0__py3-none-any.whl

seekrai/types/agents/runs.py

@@ -0,0 +1,117 @@
+import datetime
+from enum import Enum
+from typing import Any, Optional, Union
+
+from pydantic import Field
+
+from seekrai.types.abstract import BaseModel
+
+
+class RunRequest(BaseModel):
+    """Request model for creating a run."""
+
+    agent_id: str = Field(default="default_agent")
+
+
+class RunResponse(BaseModel):
+    """Response model for run creation."""
+
+    run_id: str
+    thread_id: str
+    status: str
+
+
+class RunStatus(str, Enum):
+    """Available status for a run."""
+
+    QUEUED = "queued"
+    IN_PROGRESS = "running"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELED = "canceled"
+
+
+class RunUsage(BaseModel):
+    """Aggregated usage metrics for a complete run execution.
+
+    Tracks token consumption for both prompts and completions.
+    """
+
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class RunStepUsage(BaseModel):
+    """Usage metrics for a single step within a run.
+
+    Tracks token consumption at the individual step level.
+    """
+
+    prompt_tokens: int
+    completion_tokens: int
+    total_tokens: int
+
+
+class Run(BaseModel):
+    """Represents a single execution within a thread.
+
+    A run encompasses the entire lifecycle of processing, from receiving
+    the initial prompt to delivering the final response. Runs track
+    execution status, timing, model parameters, and resource usage.
+    """
+
+    id: str
+    object: str = "thread.run"
+    created_at: datetime.datetime
+    agent_id: str
+    thread_id: str
+    status: RunStatus  # e.g. queued, in_progress, completed, failed, canceled
+    is_active: bool = False  # Indicates if this run is actively executing on its thread
+
+    started_at: Optional[datetime.datetime] = None
+    expires_at: Optional[datetime.datetime] = None
+    cancelled_at: Optional[datetime.datetime] = None
+    failed_at: Optional[datetime.datetime] = None
+    completed_at: Optional[datetime.datetime] = None
+
+    model: str
+    instructions: Optional[str] = None
+    tools: list[dict[str, Any]] = Field(default_factory=list)
+    meta_data: dict[str, Any] = Field(default_factory=dict)
+
+    usage: Optional[RunUsage] = None
+
+    temperature: float
+    top_p: float
+    max_completion_tokens: Optional[int] = None
+    truncation_strategy: dict[str, Any] = Field(default_factory=dict)
+    response_format: Union[str, dict[str, Any]] = "auto"
+    tool_choice: Union[str, dict[str, Any]] = "auto"
+    parallel_tool_calls: bool
+
+
+class RunStep(BaseModel):
+    """A single atomic operation within a run.
+
+    Steps can include message creation, tool calls, or other internal actions.
+    Each step is associated with a specific run, assistant, and thread, and
+    includes details about its type, status, and resource usage.
+    """
+
+    id: str
+    object: str = "thread.run.step"
+    created_at: datetime.datetime
+    run_id: str
+    agent_id: str
+    thread_id: str
+
+    # E.g. 'message_creation', 'tool_call', etc.
+    type: str
+
+    # E.g. 'completed', 'failed', etc.
+    status: str
+
+    completed_at: Optional[datetime.datetime] = None
+    meta_data: dict[str, Any] = Field(default_factory=dict)
+    usage: Optional[RunStepUsage] = None

seekrai/types/agents/threads.py

@@ -0,0 +1,265 @@
+import datetime
+from enum import Enum
+from typing import Any, Literal, Optional, Sequence, Union
+from uuid import uuid4
+
+from pydantic import Field
+
+from seekrai.types.abstract import BaseModel
+
+
+class ThreadCreateRequest(BaseModel):
+    """Pydantic model for creating a thread request.
+
+    Attributes:
+        meta_data (Optional[dict[str, Any]]): Optional metadata dictionary.
+    """
+
+    meta_data: Optional[dict[str, Any]] = None
+
+
+class MessageUpdateRequest(BaseModel):
+    """Pydantic model for updating a message request.
+
+    Attributes:
+        content (Optional[str]): Optional content of the message.
+        meta_data (Optional[dict[str, Any]]): Optional metadata dictionary.
+    """
+
+    content: Optional[str] = None
+    meta_data: Optional[dict[str, Any]] = None
+
+
+class ThreadStatus(str, Enum):
+    """Available status for a thread."""
+
+    AVAILABLE = "available"
+    LOCKED = "locked"
+
+
+class Thread(BaseModel):
+    """A thread is a top-level conversation container.
+
+    Threads can contain multiple messages from various assistants and users,
+    providing a complete history of an interaction sequence.
+    """
+
+    id: str
+    object: str = "thread"
+    created_at: datetime.datetime
+    status: ThreadStatus = ThreadStatus.AVAILABLE
+    active_run_id: Optional[str] = None
+    meta_data: dict[str, Any] = Field(default_factory=dict)
+
+
+class StreamReasoningChunk(BaseModel):
+    """A chunk of reasoning output from a streaming tool."""
+
+    type: Literal["stream_reasoning"]
+    reasoning: str
+    meta_data: dict[str, Any]
+
+
+class StreamTextChunk(BaseModel):
+    """A chunk of text output from a streaming tool."""
+
+    text: str
+    type: Literal["streaming_chunk"]
+    meta_data: dict[str, Any]
+
+
+class StreamingToolChunk(BaseModel):
+    """A chunk of output from a streaming tool."""
+
+    type: Literal["streaming_tool"]
+    tool_id: str
+    tool_name: str
+    tool_output: dict[str, Any]
+    meta_data: dict[str, Any]
+
+
+class StreamingToolRequest(BaseModel):
+    """A chunk of output from a streaming tool."""
+
+    type: Literal["streaming_tool_request"]
+    tool_id: str
+    tool_args: dict[str, Any]
+    meta_data: dict[str, Any]
+    tool_name: str
+
+
+class StreamingToolResponse(BaseModel):
+    """A chunk of output from a streaming tool."""
+
+    type: Literal["streaming_tool_response"]
+    tool_id: str
+    tool_output: dict[str, Any]
+    meta_data: dict[str, Any]
+
+
+class StreamNodeHeaderChunk(BaseModel):
+    """Represents a special 'header' announcement that we're at a particular Node."""
+
+    type: Literal["node_header"]
+    node_type: str  # e.g. "ModelRequestNode" or "CallToolsNode"
+    description: str  # e.g. "streaming partial request tokens"
+    meta_data: dict[str, Any]
+
+
+class StreamUserPromptChunk(BaseModel):
+    """Represents the user prompt node."""
+
+    type: Literal["user_prompt"]
+    user_prompt: Union[str, Sequence[Any]]
+    meta_data: dict[str, Any]
+
+
+class StreamPartStartEventChunk(BaseModel):
+    """Represents a PartStartEvent."""
+
+    type: Literal["part_start_event"]
+    part_index: int
+    part_content: str  # whatever event.part!r is
+    meta_data: dict[str, Any]
+
+
+class StreamToolCallPartDeltaChunk(BaseModel):
+    """Represents a ToolCallPartDelta."""
+
+    type: Literal["tool_call_part_delta"]
+    part_index: int
+    args_delta: Any  # or dict[str, Any] depending on your usage
+    meta_data: dict[str, Any]
+
+
+class StreamFinalResultEventChunk(BaseModel):
+    """Represents a FinalResultEvent."""
+
+    type: Literal["final_result_event"]
+    tool_name: Optional[str]
+    meta_data: dict[str, Any]
+
+
+class StreamEndNodeChunk(BaseModel):
+    """Represents the final agent output at the EndNode."""
+
+    type: Literal["end_node"]
+    final_output: str
+    meta_data: dict[str, Any]
+
+
+class StreamTextDeltaChunk(BaseModel):
+    """Represents partial text tokens returned by the model (TextPartDelta)."""
+
+    type: Literal["text_delta"]
+    text: str  # chunk of partial text
+    meta_data: dict[str, Any]
+
+
+StreamChunkDataTypes = Union[
+    StreamReasoningChunk,
+    StreamTextChunk,
+    StreamingToolRequest,
+    StreamingToolResponse,
+    StreamNodeHeaderChunk,
+    StreamUserPromptChunk,
+    StreamPartStartEventChunk,
+    StreamToolCallPartDeltaChunk,
+    StreamFinalResultEventChunk,
+    StreamEndNodeChunk,
+    StreamTextDeltaChunk,
+]
+
+
+class StreamChunk(BaseModel):
+    """A single chunk of streaming output from a tool."""
+
+    data: StreamChunkDataTypes
+    type: Literal["streaming_chunk"]
+    meta_data: dict[str, Any]
+
+
+class InputText(BaseModel):
+    """A text input to be sent to the model."""
+
+    text: str
+    type: Literal["input_text"]
+
+
+class InputImage(BaseModel):
+    """An image input to be sent to the model."""
+
+    detail: Literal["high", "low", "auto"]
+    """The detail level of the image to be sent to the model.
+    One of `high`, `low`, or `auto`. Defaults to `auto`.
+    """
+
+    type: Literal["input_image"]
+
+    file_id: Optional[str]
+    """The ID of the file to be sent to the model."""
+
+    image_url: Optional[str]
+    """The URL of the image to be sent to the model.
+
+    A fully qualified URL or base64 encoded image in a data URL.
+    """
+
+
+class InputFile(BaseModel):
+    """A file input to be sent to the model."""
+
+    type: Literal["input_file"]
+    """The type of the input item. Always `input_file`."""
+
+    file_id: str
+    """The ID of the file to be sent to the model."""
+
+
+InputMessage = Union[InputText, InputImage, InputFile]
+
+
+class OutputText(BaseModel):
+    """A text output from the model."""
+
+    text: str
+    type: Literal["output_text"]
+    annotations: list[str]
+
+
+class OutputGuardrail(BaseModel):
+    """A guardrail output from the model."""
+
+    type: Literal["output_guardrail"]
+    text: str
+    guardrail: list[dict[str, Any]]
+
+
+OutputMessage = Union[OutputText, OutputGuardrail]
+
+
+ThreadMessageContentType = Union[str, list[InputMessage], list[OutputMessage]]
+
+
+class ThreadMessage(BaseModel):
+    """A single piece of content within a thread.
+
+    Messages can be either requests to a model (user prompts, system prompts,
+    or tool returns) or responses from the model. Each message is associated
+    with a specific thread and optionally with an assistant and run.
+    """
+
+    id: str = Field(default_factory=lambda: str(uuid4()))
+    object: str = "thread.message"
+    thread_id: str  # FK to Thread
+    role: str  # e.g. 'user', 'assistant', 'system', 'tool'
+
+    # content can be either
+    # - a single string (for text messages)
+    # - a list of InputMessage objects (for model inputs)
+    # - a list of OutputMessage objects (for model outputs)
+    content: ThreadMessageContentType
+
+    agent_id: Optional[str] = None  # If this message was sent by an assistant
+    run_id: Optional[str] = None  # If it's part of a Run
+    meta_data: dict[str, Any] = Field(default_factory=dict)

seekrai/types/agents/tools/__init__.py

@@ -0,0 +1,16 @@
+from seekrai.types.agents.tools.env_model_config import EnvConfig
+from seekrai.types.agents.tools.schemas import FileSearch, FileSearchEnv
+from seekrai.types.agents.tools.tool import ToolBase, ToolType
+from seekrai.types.agents.tools.tool_env_types import Env
+from seekrai.types.agents.tools.tool_types import Tool
+
+
+__all__ = [
+    "ToolBase",
+    "ToolType",
+    "EnvConfig",
+    "Env",
+    "Tool",
+    "FileSearch",
+    "FileSearchEnv",
+]

seekrai/types/agents/tools/env_model_config.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel, ConfigDict
+
+
+class EnvConfig(BaseModel):
+    model_config = ConfigDict(
+        alias_generator=lambda k: k.upper(), populate_by_name=True
+    )

seekrai/types/agents/tools/schemas/__init__.py

@@ -0,0 +1,8 @@
+from seekrai.types.agents.tools.schemas.file_search import FileSearch
+from seekrai.types.agents.tools.schemas.file_search_env import FileSearchEnv
+
+
+__all__ = [
+    "FileSearch",
+    "FileSearchEnv",
+]

seekrai/types/agents/tools/schemas/file_search.py

@@ -0,0 +1,9 @@
+from typing import Literal
+
+from seekrai.types.agents.tools.schemas.file_search_env import FileSearchEnv
+from seekrai.types.agents.tools.tool import ToolBase, ToolType
+
+
+class FileSearch(ToolBase):
+    name: Literal[ToolType.FILE_SEARCH] = ToolType.FILE_SEARCH
+    tool_env: FileSearchEnv

seekrai/types/agents/tools/schemas/file_search_env.py

@@ -0,0 +1,11 @@
+from pydantic import Field
+
+from seekrai.types.agents.tools.env_model_config import EnvConfig
+
+
+# TODO: figure out better way of creating tool environment models (within tool ideally), but retaining separate model_configs
+class FileSearchEnv(EnvConfig):
+    file_search_index: str
+    document_tool_desc: str
+    top_k: int = Field(default=10)
+    score_threshold: int = Field(default=0)

seekrai/types/agents/tools/tool.py

@@ -0,0 +1,14 @@
+import enum
+
+from pydantic import BaseModel
+
+from seekrai.types.agents.tools.tool_env_types import Env
+
+
+class ToolType(str, enum.Enum):
+    FILE_SEARCH = "file_search"
+
+
+class ToolBase(BaseModel):
+    name: ToolType
+    tool_env: Env

seekrai/types/agents/tools/tool_env_types.py

@@ -0,0 +1,4 @@
+from seekrai.types.agents.tools.schemas.file_search_env import FileSearchEnv
+
+
+Env = FileSearchEnv  # will be a Union of tool envs when more are added

seekrai/types/agents/tools/tool_types.py

@@ -0,0 +1,10 @@
+from typing import Annotated
+
+from pydantic import Field
+
+from seekrai.types.agents.tools.schemas.file_search import FileSearch
+
+
+Tool = Annotated[
+    FileSearch, Field(discriminator="name")
+]  # will be a Union of tools when more are added

seekrai/types/alignment.py

@@ -12,7 +12,8 @@ from seekrai.types.abstract import BaseModel
 class AlignmentType(str, Enum):
     PRINCIPLE = "principle"
     CHAIN_OF_THOUGHT = "chain_of_thought"
-    RAFT = "raft"
+    RAFT = "raft"  # deprecated - use CONTEXT_DATA instead
+    CONTEXT_DATA = "context_data"
 
 
 class AlignmentRequest(BaseModel):
@@ -24,7 +25,7 @@ class AlignmentRequest(BaseModel):
     )
     type: AlignmentType = Field(
         default=AlignmentType.PRINCIPLE,
-        description="Type of alignment task (principle, chain_of_thought, or raft)",
+        description="Type of alignment task (principle, chain_of_thought, or context_data)",
     )
 
 
@@ -59,7 +60,10 @@ class AlignmentJobStatus(str, Enum):
 class AlignmentResponse(BaseModel):
     id: Optional[str] = Field(default=..., description="Alignment job ID")
     created_at: datetime | None = None
+    updated_at: datetime | None = None
     status: AlignmentJobStatus | None = None
+    current_step: str | None = None
+    progress: str | None = None
 
 
 class AlignmentList(BaseModel):

seekrai/types/files.py

@@ -14,6 +14,8 @@ class FilePurpose(str, Enum):
     FineTune = "fine-tune"
     PreTrain = "pre-train"
     Alignment = "alignment"
+    Ingestion = "ingestion"
+    VectorIngestion = "vector-ingestion"
 
 
 class TrainingFileType(str, Enum):
@@ -94,6 +96,7 @@ class FileResponse(BaseModel):
     # file byte size
     bytes: int | None = None
     created_by: str | None = None  # TODO - fix this later
+    deleted: bool | None = None
 
 
 class FileList(BaseModel):

seekrai/types/finetune.py

@@ -113,6 +113,7 @@ class AcceleratorType(str, Enum):
     A100 = "A100"
     A10 = "A10"
     H100 = "H100"
+    MI300X = "MI300X"
 
 
 class InfrastructureConfig(BaseModel):

seekrai/types/models.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import datetime
 from enum import Enum
 from typing import List, Literal
 
@@ -29,10 +30,12 @@ class ModelResponse(BaseModel):
     # object type
     object: Literal[ObjectType.Model]
     created: int | None = None
+    created_at: datetime | None = None
     # model type
     type: ModelType | None = None
     name: str | None = None
     bytes: int | None = None
+    model_type: str | None = None
     # # model creator organization
     # organization: str | None = None
     # # link to model resource

seekrai/types/vectordb.py

@@ -0,0 +1,78 @@
+from datetime import datetime
+from typing import List, Literal, Optional
+
+from pydantic import BaseModel, Field
+
+
+class VectorDatabaseCreate(BaseModel):
+    """Request model for creating a new vector database."""
+
+    name: str = Field(..., description="Name of the vector database")
+    model: str = Field(..., description="Model used to generate the vectors")
+    description: Optional[str] = Field(None, description="Optional description")
+
+
+class VectorDatabaseResponse(BaseModel):
+    """Response model for a vector database."""
+
+    id: str
+    name: str
+    model: str
+    dimension: int
+    description: Optional[str]
+    created_at: datetime
+    updated_at: datetime
+    file_count: int
+    size_in_bytes: Optional[int] = None
+
+
+class VectorDatabaseList(BaseModel):
+    """Response model for a list of vector databases."""
+
+    object: Literal["list"]
+    data: List[VectorDatabaseResponse]
+
+
+class VectorDatabaseIngestionRequest(BaseModel):
+    """Request model for creating a new vector database ingestion job."""
+
+    file_ids: List[str] = Field(..., description="List of file IDs to ingest")
+    method: str = Field(..., description="Method to use for ingestion")
+    token_count: int = Field(default=800, description="Token count for ingestion")
+    overlap_tokens: int = Field(default=100, description="Overlap tokens for ingestion")
+
+
+class VectorDatabaseIngestionResponse(BaseModel):
+    """Response model for a vector database ingestion job."""
+
+    id: str
+    vector_database_id: str
+    status: str
+    created_at: datetime
+    updated_at: datetime
+    error_message: Optional[str]
+    file_ids: List[str]
+    metaflow_run_id: Optional[str]
+
+
+class VectorDatabaseIngestionList(BaseModel):
+    """Response model for a list of vector database ingestion jobs."""
+
+    object: Literal["list"]
+    data: List[VectorDatabaseIngestionResponse]
+
+
+class VectorDatabaseFileResponse(BaseModel):
+    """Response model for a vector database file."""
+
+    id: str
+    vector_database_id: str
+    filename: str
+    created_at: datetime
+
+
+class VectorDatabaseFileList(BaseModel):
+    """Response model for a list of vector database files."""
+
+    object: Literal["list"]
+    data: List[VectorDatabaseFileResponse]

{seekrai-0.4.4.dist-info → seekrai-0.5.0.dist-info}/METADATA

@@ -1,8 +1,7 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: seekrai
-Version: 0.4.4
+Version: 0.5.0
 Summary: Python client for SeekrAI
-Home-page: https://gitlab.cb.ntent.com/ml/seekr-py
 License: Apache-2.0
 Author: SeekrFlow
 Author-email: support@seekr.com
@@ -14,6 +13,7 @@ Classifier: Programming Language :: Python :: 3.9
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
 Requires-Dist: click (>=8.1.7,<9.0.0)
 Requires-Dist: eval-type-backport (>=0.1.3,<0.3.0)
 Requires-Dist: filelock (>=3.13.1,<4.0.0)