retab 0.0.38__py3-none-any.whl → 0.0.39__py3-none-any.whl

retab/_utils/ai_models.py

@@ -1,74 +1,121 @@
+import os
+import yaml
 from typing import get_args
 
-from ..types.ai_models import AIProvider, GeminiModel, OpenAIModel, xAI_Model
+from ..types.ai_models import AIProvider, GeminiModel, OpenAIModel, xAI_Model, RetabModel, PureLLMModel, ModelCard
 
+MODEL_CARDS_DIR = os.path.join(os.path.dirname(__file__), "_model_cards")
 
-def find_provider_from_model(model: str) -> AIProvider:
-    if model in get_args(OpenAIModel):
-        return "OpenAI"
-    elif ":" in model:
-        # Handle fine-tuned models
-        ft, base_model, model_id = model.split(":", 2)
-        if base_model in get_args(OpenAIModel):
-            return "OpenAI"
-    # elif model in get_args(AnthropicModel):
-    #    return "Anthropic"
-    elif model in get_args(xAI_Model):
-        return "xAI"
-    elif model in get_args(GeminiModel):
-        return "Gemini"
-    raise ValueError(f"Could not determine AI provider for model: {model}")
+def merge_model_cards(base: dict, override: dict) -> dict:
+    result = base.copy()
+    for key, value in override.items():
+        if key == "inherits":
+            continue
+        if isinstance(value, dict) and key in result:
+            result[key] = merge_model_cards(result[key], value)
+        else:
+            result[key] = value
+    return result
 
+def load_model_cards(yaml_file: str) -> list[ModelCard]:
+    raw_cards = yaml.safe_load(open(yaml_file))
+    name_to_card = {c["model"]: c for c in raw_cards if "inherits" not in c}
 
-def assert_valid_model_extraction(model: str) -> None:
-    if model in get_args(OpenAIModel):
-        return
-    elif ":" in model:
-        # Handle fine-tuned models
-        ft, base_model, model_id = model.split(":", 2)
-        if base_model in get_args(OpenAIModel):
-            return
-    # elif model in get_args(AnthropicModel):
-    #    return
-    elif model in get_args(xAI_Model):
-        return
-    elif model in get_args(GeminiModel):
-        return
-    raise ValueError(
-        f"Invalid model for extraction: {model}.\nValid OpenAI models: {get_args(OpenAIModel)}\n"
-        # f"Valid Anthropic models: {get_args(AnthropicModel)}\n"
-        # f"Valid xAI models: {get_args(xAI_Model)}\n"
-        # f"Valid Gemini models: {get_args(GeminiModel)}"
-    )
+    final_cards = []
+    for card in raw_cards:
+        if "inherits" in card:
+            parent = name_to_card[card["inherits"]]
+            merged = merge_model_cards(parent, card)
+            final_cards.append(ModelCard(**merged))
+        else:
+            final_cards.append(ModelCard(**card))
+    return final_cards
 
+# Load all model cards
+model_cards = sum([
+    load_model_cards(os.path.join(MODEL_CARDS_DIR, "openai.yaml")),
+    load_model_cards(os.path.join(MODEL_CARDS_DIR, "anthropic.yaml")),
+    load_model_cards(os.path.join(MODEL_CARDS_DIR, "xai.yaml")),
+    load_model_cards(os.path.join(MODEL_CARDS_DIR, "gemini.yaml")),
+    load_model_cards(os.path.join(MODEL_CARDS_DIR, "auto.yaml")),
+], [])
+model_cards_dict = {card.model: card for card in model_cards}
 
-def assert_valid_model_batch_processing(model: str) -> None:
-    """Assert that the model is either a standard OpenAI model or a valid fine-tuned model.
 
-    Valid formats:
-    - Standard model: Must be in OpenAIModel
-    - Fine-tuned model: Must be {base_model}:{id} where base_model is in OpenAIModel
+# Validate that model cards
+all_model_names = set(model_cards_dict.keys())
+if all_model_names.symmetric_difference(set(get_args(PureLLMModel))):
+    raise ValueError(f"Mismatch between model cards and PureLLMModel type: {all_model_names.symmetric_difference(set(get_args(PureLLMModel)))}")
+
+
+def get_model_from_model_id(model_id: str) -> str:
+    """
+    Get the model name from the model id.
+    """
+    if model_id.startswith("ft:"):
+        parts = model_id.split(":")
+        return parts[1]
+    else:
+        return model_id
+
+
+def get_model_card(model: str) -> ModelCard:
+    """
+    Get the model card for a specific model.
+
+    Args:
+        model: The model name to look up
+
+    Returns:
+        The ModelCard for the specified model
 
     Raises:
-        ValueError: If the model format is invalid
+        ValueError: If no model card is found for the specified model
     """
-    if model in get_args(OpenAIModel):
-        return
+    model_name = get_model_from_model_id(model)
+    if model_name in model_cards_dict:
+        model_card = ModelCard(**model_cards_dict[model_name].model_dump())
+        if model_name != model:
+            # Fine-tuned model -> Change the name
+            model_card.model = model
+            # Remove the fine-tuning feature (if exists)
+            try:
+                model_card.capabilities.features.remove("fine_tuning")
+            except ValueError:
+                pass
+        return model_card
+
+    raise ValueError(f"No model card found for model: {model_name}")
 
-    try:
-        ft, base_model, model_id = model.split(":", 2)
-        if base_model not in get_args(OpenAIModel):
-            raise ValueError(f"Invalid base model in fine-tuned model '{model}'. Base model must be one of: {get_args(OpenAIModel)}")
-        if not model_id or not model_id.strip():
-            raise ValueError(f"Model ID cannot be empty in fine-tuned model '{model}'")
+
+def get_provider_for_model(model_id: str) -> AIProvider:
+    """
+    Determine the AI provider associated with the given model identifier.
+    Returns one of: "Anthropic", "xAI", "OpenAI", "Gemini", "Retab" or None if unknown.
+    """
+    model_name = get_model_from_model_id(model_id)
+    # if model_name in get_args(AnthropicModel):
+    #     return "Anthropic"
+    # if model_name in get_args(xAI_Model):
+    #     return "xAI"
+    if model_name in get_args(OpenAIModel):
+        return "OpenAI"
+    if model_name in get_args(GeminiModel):
+        return "Gemini"
+    if model_name in get_args(RetabModel):
+        return "Retab"
+    raise ValueError(f"Unknown model: {model_name}")
+
+
+def assert_valid_model_extraction(model: str) -> None:
+    try: 
+        get_provider_for_model(model)
     except ValueError:
-        if ":" not in model:
-            raise ValueError(
-                f"Invalid model format: {model}. Must be either:\n"
-                f"1. A standard model: {get_args(OpenAIModel)}\n"
-                f"2. A fine-tuned model in format 'base_model:id' where base_model is one of the standard models"
-            ) from None
-        raise
+        raise ValueError(
+            f"Invalid model for extraction: {model}.\nValid OpenAI models: {get_args(OpenAIModel)}\n"
+            f"Valid xAI models: {get_args(xAI_Model)}\n"
+            f"Valid Gemini models: {get_args(GeminiModel)}"
+        ) from None
 
 
 def assert_valid_model_schema_generation(model: str) -> None:
@@ -81,20 +128,11 @@ def assert_valid_model_schema_generation(model: str) -> None:
     Raises:
         ValueError: If the model format is invalid
     """
-    if model in get_args(OpenAIModel):
+    if get_model_from_model_id(model) in get_args(OpenAIModel):
         return
-
-    try:
-        ft, base_model, model_id = model.split(":", 2)
-        if base_model not in get_args(OpenAIModel):
-            raise ValueError(f"Invalid base model in fine-tuned model '{model}'. Base model must be one of: {get_args(OpenAIModel)}")
-        if not model_id or not model_id.strip():
-            raise ValueError(f"Model ID cannot be empty in fine-tuned model '{model}'")
-    except ValueError:
-        if ":" not in model:
-            raise ValueError(
+    else:
+        raise ValueError(
                 f"Invalid model format: {model}. Must be either:\n"
                 f"1. A standard model: {get_args(OpenAIModel)}\n"
-                f"2. A fine-tuned model in format 'base_model:id' where base_model is one of the standard models"
-            ) from None
-        raise
+                f"2. A fine-tuned model in format 'base_model:id' where base_model is one of the standard openai models"
+            ) from None 

retab/_utils/responses.py

@@ -17,7 +17,7 @@ from openai.types.responses.response_input_param import ResponseInputItemParam
 from openai.types.responses.response_input_text_param import ResponseInputTextParam
 
 from ..types.chat import ChatCompletionRetabMessage
-from ..types.documents.extractions import UiParsedChatCompletion, UiParsedChoice
+from ..types.documents.extractions import RetabParsedChatCompletion, RetabParsedChoice
 
 
 def convert_to_openai_format(messages: list[ChatCompletionRetabMessage]) -> list[ResponseInputItemParam]:
@@ -116,17 +116,17 @@ def convert_from_openai_format(messages: list[ResponseInputItemParam]) -> list[C
     return formatted_messages
 
 
-def parse_openai_responses_response(response: Response) -> UiParsedChatCompletion:
+def parse_openai_responses_response(response: Response) -> RetabParsedChatCompletion:
     """
-    Convert an OpenAI Response (Responses API) to UiParsedChatCompletion type.
+    Convert an OpenAI Response (Responses API) to RetabParsedChatCompletion type.
 
     Args:
         response: Response from OpenAI Responses API
 
     Returns:
-        Parsed response in UiParsedChatCompletion format
+        Parsed response in RetabParsedChatCompletion format
     """
-    # Create the UiParsedChatCompletion object
+    # Create the RetabParsedChatCompletion object
     if response.usage:
         usage = CompletionUsage(
             prompt_tokens=response.usage.input_tokens,
@@ -148,7 +148,7 @@ def parse_openai_responses_response(response: Response) -> UiParsedChatCompletio
     result_object = from_json(bytes(output_text, "utf-8"), partial_mode=True)  # Attempt to parse the result even if EOF is reached
 
     choices.append(
-        UiParsedChoice(
+        RetabParsedChoice(
             index=0,
             message=ParsedChatCompletionMessage(
                 role="assistant",
@@ -158,7 +158,7 @@ def parse_openai_responses_response(response: Response) -> UiParsedChatCompletio
         )
     )
 
-    return UiParsedChatCompletion(
+    return RetabParsedChatCompletion(
         id=response.id,
         choices=choices,
         created=int(datetime.datetime.now().timestamp()),

retab/_utils/usage/usage.py

@@ -4,7 +4,8 @@ from openai.types.completion_usage import CompletionUsage
 from pydantic import BaseModel, Field
 
 # https://platform.openai.com/docs/guides/prompt-caching
-from ...types.ai_models import Amount, Pricing, get_model_card
+from ...types.ai_models import Amount, Pricing
+from ..._utils.ai_models import get_model_card
 
 # ─── PRICING MODELS ────────────────────────────────────────────────────────────
 

retab/resources/consensus/completions.py

@@ -7,8 +7,8 @@ from pydantic import BaseModel as ResponseFormatT
 from ..._resource import AsyncAPIResource, SyncAPIResource
 from ..._utils.ai_models import assert_valid_model_extraction
 from ...types.chat import ChatCompletionRetabMessage
-from ...types.completions import UiChatCompletionsRequest
-from ...types.documents.extractions import UiParsedChatCompletion
+from ...types.completions import RetabChatCompletionsRequest
+from ...types.documents.extractions import RetabParsedChatCompletion
 from ...types.schemas.object import Schema
 from ...types.standards import PreparedRequest
 
@@ -31,7 +31,7 @@ class BaseCompletionsMixin:
 
         schema_obj = Schema(json_schema=json_schema)
 
-        request = UiChatCompletionsRequest(
+        request = RetabChatCompletionsRequest(
             model=model,
             messages=messages,
             response_format={
@@ -66,7 +66,7 @@ class BaseCompletionsMixin:
 
         schema_obj = Schema(json_schema=json_schema)
 
-        request = UiChatCompletionsRequest(
+        request = RetabChatCompletionsRequest(
             model=model,
             messages=messages,
             response_format={
@@ -98,7 +98,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
         n_consensus: int = 1,
         idempotency_key: str | None = None,
         stream: bool = False,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """
         Create a completion using the Retab API.
         """
@@ -116,7 +116,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
 
         response = self._client._prepared_request(request)
 
-        return UiParsedChatCompletion.model_validate(response)
+        return RetabParsedChatCompletion.model_validate(response)
 
     def parse(
         self,
@@ -127,7 +127,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
         reasoning_effort: ChatCompletionReasoningEffort = "medium",
         n_consensus: int = 1,
         idempotency_key: str | None = None,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """
         Parse messages using the Retab API to extract structured data according to the provided JSON schema.
 
@@ -141,7 +141,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
             store: Whether to store the data in the Retab database
 
         Returns:
-            UiParsedChatCompletion: Parsed response from the API
+            RetabParsedChatCompletion: Parsed response from the API
         """
         request = self.prepare_parse(
             response_format=response_format,
@@ -155,7 +155,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
         )
         response = self._client._prepared_request(request)
 
-        return UiParsedChatCompletion.model_validate(response)
+        return RetabParsedChatCompletion.model_validate(response)
 
 
 class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
@@ -171,7 +171,7 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
         n_consensus: int = 1,
         idempotency_key: str | None = None,
         stream: bool = False,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """
         Create a completion using the Retab API.
         """
@@ -188,7 +188,7 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
         )
 
         response = await self._client._prepared_request(request)
-        return UiParsedChatCompletion.model_validate(response)
+        return RetabParsedChatCompletion.model_validate(response)
 
     async def parse(
         self,
@@ -199,7 +199,7 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
         reasoning_effort: ChatCompletionReasoningEffort = "medium",
         n_consensus: int = 1,
         idempotency_key: str | None = None,
-    ) -> UiParsedChatCompletion:
+    ) -> RetabParsedChatCompletion:
         """
         Parse messages using the Retab API asynchronously.
 
@@ -213,7 +213,7 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
             idempotency_key: Idempotency key for request
 
         Returns:
-            UiParsedChatCompletion: Parsed response from the API
+            RetabParsedChatCompletion: Parsed response from the API
         """
         request = self.prepare_parse(
             response_format=response_format,
@@ -226,4 +226,4 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
             idempotency_key=idempotency_key,
         )
         response = await self._client._prepared_request(request)
-        return UiParsedChatCompletion.model_validate(response)
+        return RetabParsedChatCompletion.model_validate(response)

retab/resources/consensus/completions_stream.py

@@ -13,8 +13,8 @@ from ..._utils.ai_models import assert_valid_model_extraction
 from ..._utils.json_schema import unflatten_dict
 from ..._utils.stream_context_managers import as_async_context_manager, as_context_manager
 from ...types.chat import ChatCompletionRetabMessage
-from ...types.completions import UiChatCompletionsRequest
-from ...types.documents.extractions import UiParsedChatCompletion, UiParsedChatCompletionChunk, UiParsedChoice
+from ...types.completions import RetabChatCompletionsRequest
+from ...types.documents.extractions import RetabParsedChatCompletion, RetabParsedChatCompletionChunk, RetabParsedChoice
 from ...types.schemas.object import Schema
 from ...types.standards import PreparedRequest
 
@@ -36,7 +36,7 @@ class BaseCompletionsMixin:
         json_schema = response_format.model_json_schema()
         schema_obj = Schema(json_schema=json_schema)
 
-        request = UiChatCompletionsRequest(
+        request = RetabChatCompletionsRequest(
             messages=messages,
             response_format={
                 "type": "json_schema",
@@ -73,7 +73,7 @@ class BaseCompletionsMixin:
         schema_obj = Schema(json_schema=json_schema)
 
         # Validate DocumentAPIRequest data (raises exception if invalid)
-        request = UiChatCompletionsRequest(
+        request = RetabChatCompletionsRequest(
             model=model,
             messages=messages,
             response_format={
@@ -106,7 +106,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
         reasoning_effort: ChatCompletionReasoningEffort = "medium",
         n_consensus: int = 1,
         idempotency_key: str | None = None,
-    ) -> Generator[UiParsedChatCompletion, None, None]:
+    ) -> Generator[RetabParsedChatCompletion, None, None]:
         """
         Process messages using the Retab API with streaming enabled.
 
@@ -119,7 +119,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
             idempotency_key: Idempotency key for request
 
         Returns:
-            Generator[UiParsedChatCompletion]: Stream of parsed responses
+            Generator[RetabParsedChatCompletion]: Stream of parsed responses
 
         Usage:
         ```python
@@ -140,16 +140,16 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
         )
 
         # Request the stream and return a context manager
-        ui_parsed_chat_completion_cum_chunk: UiParsedChatCompletionChunk | None = None
-        # Initialize the UiParsedChatCompletion object
-        ui_parsed_completion: UiParsedChatCompletion = UiParsedChatCompletion(
+        ui_parsed_chat_completion_cum_chunk: RetabParsedChatCompletionChunk | None = None
+        # Initialize the RetabParsedChatCompletion object
+        ui_parsed_completion: RetabParsedChatCompletion = RetabParsedChatCompletion(
             id="",
             created=0,
             model="",
             object="chat.completion",
             likelihoods={},
             choices=[
-                UiParsedChoice(
+                RetabParsedChoice(
                     index=0,
                     message=ParsedChatCompletionMessage(content="", role="assistant"),
                     finish_reason=None,
@@ -160,7 +160,7 @@ class Completions(SyncAPIResource, BaseCompletionsMixin):
         for chunk_json in self._client._prepared_request_stream(request):
             if not chunk_json:
                 continue
-            ui_parsed_chat_completion_cum_chunk = UiParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
+            ui_parsed_chat_completion_cum_chunk = RetabParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
             # Basic stuff
             ui_parsed_completion.id = ui_parsed_chat_completion_cum_chunk.id
             ui_parsed_completion.created = ui_parsed_chat_completion_cum_chunk.created
@@ -192,7 +192,7 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
         reasoning_effort: ChatCompletionReasoningEffort = "medium",
         n_consensus: int = 1,
         idempotency_key: str | None = None,
-    ) -> AsyncGenerator[UiParsedChatCompletion, None]:
+    ) -> AsyncGenerator[RetabParsedChatCompletion, None]:
         """
         Parse messages using the Retab API asynchronously with streaming.
 
@@ -206,7 +206,7 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
             idempotency_key: Idempotency key for request
 
         Returns:
-            AsyncGenerator[UiParsedChatCompletion]: Stream of parsed responses
+            AsyncGenerator[RetabParsedChatCompletion]: Stream of parsed responses
 
         Usage:
         ```python
@@ -227,16 +227,16 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
         )
 
         # Request the stream and return a context manager
-        ui_parsed_chat_completion_cum_chunk: UiParsedChatCompletionChunk | None = None
-        # Initialize the UiParsedChatCompletion object
-        ui_parsed_completion: UiParsedChatCompletion = UiParsedChatCompletion(
+        ui_parsed_chat_completion_cum_chunk: RetabParsedChatCompletionChunk | None = None
+        # Initialize the RetabParsedChatCompletion object
+        ui_parsed_completion: RetabParsedChatCompletion = RetabParsedChatCompletion(
             id="",
             created=0,
             model="",
             object="chat.completion",
             likelihoods={},
             choices=[
-                UiParsedChoice(
+                RetabParsedChoice(
                     index=0,
                     message=ParsedChatCompletionMessage(content="", role="assistant"),
                     finish_reason=None,
@@ -247,7 +247,7 @@ class AsyncCompletions(AsyncAPIResource, BaseCompletionsMixin):
         async for chunk_json in self._client._prepared_request_stream(request):
             if not chunk_json:
                 continue
-            ui_parsed_chat_completion_cum_chunk = UiParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
+            ui_parsed_chat_completion_cum_chunk = RetabParsedChatCompletionChunk.model_validate(chunk_json).chunk_accumulator(ui_parsed_chat_completion_cum_chunk)
             # Basic stuff
             ui_parsed_completion.id = ui_parsed_chat_completion_cum_chunk.id
             ui_parsed_completion.created = ui_parsed_chat_completion_cum_chunk.created

retab/resources/consensus/responses.py

@@ -8,7 +8,7 @@ from pydantic import BaseModel
 
 from ..._resource import AsyncAPIResource, SyncAPIResource
 from ..._utils.ai_models import assert_valid_model_extraction
-from ...types.completions import UiChatResponseCreateRequest
+from ...types.completions import RetabChatResponseCreateRequest
 from ...types.documents.extractions import UiResponse
 from ...types.schemas.object import Schema
 from ...types.standards import PreparedRequest
@@ -43,8 +43,8 @@ class BaseResponsesMixin:
         if instructions is None:
             instructions = schema_obj.developer_system_prompt
 
-        # Create the request object based on the UiChatResponseCreateRequest model
-        request = UiChatResponseCreateRequest(
+        # Create the request object based on the RetabChatResponseCreateRequest model
+        request = RetabChatResponseCreateRequest(
             model=model,
             input=input,
             temperature=temperature,
@@ -79,8 +79,8 @@ class BaseResponsesMixin:
         if instructions is None:
             instructions = schema_obj.developer_system_prompt
 
-        # Create the request object based on the UiChatResponseCreateRequest model
-        request = UiChatResponseCreateRequest(
+        # Create the request object based on the RetabChatResponseCreateRequest model
+        request = RetabChatResponseCreateRequest(
             model=model,
             input=input,
             temperature=temperature,

retab/resources/consensus/responses_stream.py

@@ -8,7 +8,7 @@ from pydantic import BaseModel
 from ..._resource import AsyncAPIResource, SyncAPIResource
 from ..._utils.ai_models import assert_valid_model_extraction
 from ..._utils.stream_context_managers import as_async_context_manager, as_context_manager
-from ...types.completions import UiChatResponseCreateRequest
+from ...types.completions import RetabChatResponseCreateRequest
 from ...types.documents.extractions import UiResponse
 from ...types.schemas.object import Schema
 from ...types.standards import PreparedRequest
@@ -44,8 +44,8 @@ class BaseResponsesMixin:
         if instructions is None:
             instructions = schema_obj.developer_system_prompt
 
-        # Create the request object based on the UiChatResponseCreateRequest model
-        request = UiChatResponseCreateRequest(
+        # Create the request object based on the RetabChatResponseCreateRequest model
+        request = RetabChatResponseCreateRequest(
             model=model,
             input=input,
             temperature=temperature,
@@ -81,8 +81,8 @@ class BaseResponsesMixin:
         if instructions is None:
             instructions = schema_obj.developer_system_prompt
 
-        # Create the request object based on the UiChatResponseCreateRequest model
-        request = UiChatResponseCreateRequest(
+        # Create the request object based on the RetabChatResponseCreateRequest model
+        request = RetabChatResponseCreateRequest(
             model=model,
             input=input,
             temperature=temperature,