openaivec 0.14.4__tar.gz → 0.14.6__tar.gz

{openaivec-0.14.4 → openaivec-0.14.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: openaivec
-Version: 0.14.4
+Version: 0.14.6
 Summary: Generative mutation for tabular calculation
 Project-URL: Homepage, https://microsoft.github.io/openaivec/
 Project-URL: Repository, https://github.com/microsoft/openaivec

{openaivec-0.14.4 → openaivec-0.14.6}/src/openaivec/_provider.py

@@ -13,6 +13,7 @@ from openaivec._model import (
     OpenAIAPIKey,
     ResponsesModelName,
 )
+from openaivec._schema import SchemaInferer
 from openaivec._util import TextChunker
 
 __all__ = []
@@ -142,6 +143,13 @@ CONTAINER.register(OpenAI, provide_openai_client)
 CONTAINER.register(AsyncOpenAI, provide_async_openai_client)
 CONTAINER.register(tiktoken.Encoding, lambda: tiktoken.get_encoding("o200k_base"))
 CONTAINER.register(TextChunker, lambda: TextChunker(CONTAINER.resolve(tiktoken.Encoding)))
+CONTAINER.register(
+    SchemaInferer,
+    lambda: SchemaInferer(
+        client=CONTAINER.resolve(OpenAI),
+        model_name=CONTAINER.resolve(ResponsesModelName).value,
+    ),
+)
 
 
 def reset_environment_registrations():
@@ -160,3 +168,10 @@ def reset_environment_registrations():
     )
     CONTAINER.register(OpenAI, provide_openai_client)
     CONTAINER.register(AsyncOpenAI, provide_async_openai_client)
+    CONTAINER.register(
+        SchemaInferer,
+        lambda: SchemaInferer(
+            client=CONTAINER.resolve(OpenAI),
+            model_name=CONTAINER.resolve(ResponsesModelName).value,
+        ),
+    )

{openaivec-0.14.4 → openaivec-0.14.6}/src/openaivec/pandas_ext.py

@@ -49,6 +49,8 @@ import pandas as pd
 import tiktoken
 from openai import AsyncOpenAI, OpenAI
 
+from openaivec._schema import InferredSchema, SchemaInferenceInput, SchemaInferer
+
 __all__ = [
     "embeddings_model",
     "responses_model",
@@ -434,6 +436,61 @@ class OpenAIVecSeriesAccessor:
             **api_kwargs,
         )
 
+    def infer_schema(self, purpose: str, max_examples: int = 100) -> InferredSchema:
+        """Infer a structured data schema from Series content using AI.
+
+        This method analyzes a sample of the Series values to automatically infer
+        a structured schema that can be used for consistent data extraction.
+        The inferred schema includes field names, types, descriptions, and
+        potential enum values based on patterns found in the data.
+
+        Args:
+            purpose (str): Plain language description of how the extracted
+                structured data will be used (e.g., "Extract customer sentiment
+                signals for analytics", "Parse product features for search").
+                This guides field relevance and helps exclude irrelevant information.
+            max_examples (int): Maximum number of examples to analyze from the
+                Series. The method will sample randomly from the Series up to this
+                limit. Defaults to 100.
+
+        Returns:
+            InferredSchema: An object containing:
+                - purpose: Normalized statement of the extraction objective
+                - fields: List of field specifications with names, types, and descriptions
+                - inference_prompt: Reusable prompt for future extractions
+                - model: Dynamically generated Pydantic model for parsing
+                - task: PreparedTask for batch extraction operations
+
+        Example:
+            ```python
+            reviews = pd.Series([
+                "Great product! Fast shipping and excellent quality.",
+                "Terrible experience. Item broke after 2 days.",
+                "Average product. Price is fair but nothing special."
+            ])
+
+            # Infer schema for sentiment analysis
+            schema = reviews.ai.infer_schema(
+                purpose="Extract sentiment and product quality indicators"
+            )
+
+            # Use the inferred schema for batch extraction
+            extracted = reviews.ai.task(schema.task)
+            ```
+
+        Note:
+            The schema inference uses AI to analyze patterns in the data and may
+            require multiple attempts to produce a valid schema. Fields are limited
+            to primitive types (string, integer, float, boolean) with optional
+            enum values for categorical fields.
+        """
+        inferer = CONTAINER.resolve(SchemaInferer)
+
+        input: SchemaInferenceInput = SchemaInferenceInput(
+            examples=self._obj.sample(n=min(max_examples, len(self._obj))).tolist(), purpose=purpose
+        )
+        return inferer.infer_schema(input)
+
     def count_tokens(self) -> pd.Series:
         """Count `tiktoken` tokens per row.
 
@@ -480,6 +537,90 @@ class OpenAIVecSeriesAccessor:
             extracted.columns = [f"{self._obj.name}_{col}" for col in extracted.columns]
         return extracted
 
+    def auto_extract(
+        self,
+        purpose: str,
+        max_examples: int = 100,
+        batch_size: int | None = None,
+        show_progress: bool = False,
+        **api_kwargs,
+    ) -> pd.DataFrame:
+        """Automatically infer schema and extract structured data in one step.
+
+        This convenience method combines schema inference and data extraction into
+        a single operation. It first analyzes a sample of the Series to infer an
+        appropriate schema based on the stated purpose, then immediately applies
+        that schema to extract structured data from all values in the Series.
+
+        Args:
+            purpose (str): Plain language description of what information to extract
+                and how it will be used (e.g., "Extract product features for search",
+                "Parse customer feedback for sentiment analysis"). This guides both
+                schema inference and field selection.
+            max_examples (int): Maximum number of examples to use for schema inference.
+                A larger sample may produce more accurate schemas but increases
+                inference time. Defaults to 100.
+            batch_size (int | None): Number of requests to process in parallel during
+                extraction. Defaults to None (automatic optimization). Set to a specific
+                value to control API usage and performance.
+            show_progress (bool): Whether to display a progress bar during extraction.
+                Useful for large datasets. Defaults to False.
+            **api_kwargs: Additional OpenAI API parameters (e.g., `temperature`, `top_p`,
+                `frequency_penalty`, `presence_penalty`, `seed`) forwarded to the task execution.
+
+        Returns:
+            pd.DataFrame: A DataFrame with extracted structured data. Each inferred
+                field becomes a column, with the same index as the original Series.
+                Column names and types are determined by the inferred schema.
+
+        Example:
+            ```python
+            # Extract structured data from product reviews
+            reviews = pd.Series([
+                "Great laptop! 16GB RAM, fast SSD, battery lasts 10 hours",
+                "Decent phone. 128GB storage, camera is okay, screen is bright",
+                "Gaming desktop with RTX 4090, 32GB RAM, runs everything smoothly"
+            ])
+
+            # One-step extraction
+            extracted = reviews.ai.auto_extract(
+                purpose="Extract product specifications and performance metrics",
+                show_progress=True
+            )
+            # Result: DataFrame with columns like 'ram', 'storage', 'battery_life', etc.
+
+            # Extract sentiment and issues from support tickets
+            tickets = pd.Series([
+                "Account locked, can't reset password, very frustrated",
+                "Billing error, charged twice for subscription",
+                "Great support! Issue resolved quickly"
+            ])
+
+            features = tickets.ai.auto_extract(
+                purpose="Extract issue type and customer sentiment for support analytics"
+            )
+            ```
+
+        Note:
+            This method is ideal for exploratory data analysis when you don't have
+            a predefined schema. For production use cases with stable schemas,
+            consider using `infer_schema()` once and reusing the schema with `task()`.
+            The inferred schema is not returned, so if you need to inspect or save it,
+            use `infer_schema()` and `task()` separately.
+        """
+        schema = self._obj.ai.infer_schema(purpose=purpose, max_examples=max_examples)
+
+        return pd.DataFrame(
+            {
+                "inferred": self._obj.ai.task(
+                    task=schema.task,
+                    batch_size=batch_size,
+                    show_progress=show_progress,
+                    **api_kwargs,
+                ),
+            }
+        ).ai.extract("inferred")
+
 
 @pd.api.extensions.register_dataframe_accessor("ai")
 class OpenAIVecDataFrameAccessor:
@@ -680,6 +821,62 @@ class OpenAIVecDataFrameAccessor:
             **api_kwargs,
         )
 
+    def infer_schema(self, purpose: str, max_examples: int = 100) -> InferredSchema:
+        """Infer a structured data schema from DataFrame rows using AI.
+
+        This method analyzes a sample of DataFrame rows to automatically infer
+        a structured schema that can be used for consistent data extraction.
+        Each row is converted to JSON format and analyzed to identify patterns,
+        field types, and potential categorical values.
+
+        Args:
+            purpose (str): Plain language description of how the extracted
+                structured data will be used (e.g., "Extract operational metrics
+                for dashboard", "Parse customer attributes for segmentation").
+                This guides field relevance and helps exclude irrelevant information.
+            max_examples (int): Maximum number of rows to analyze from the
+                DataFrame. The method will sample randomly up to this limit.
+                Defaults to 100.
+
+        Returns:
+            InferredSchema: An object containing:
+                - purpose: Normalized statement of the extraction objective
+                - fields: List of field specifications with names, types, and descriptions
+                - inference_prompt: Reusable prompt for future extractions
+                - model: Dynamically generated Pydantic model for parsing
+                - task: PreparedTask for batch extraction operations
+
+        Example:
+            ```python
+            df = pd.DataFrame({
+                'text': [
+                    "Order #123: Shipped to NYC, arriving Tuesday",
+                    "Order #456: Delayed due to weather, new ETA Friday",
+                    "Order #789: Delivered to customer in LA"
+                ],
+                'timestamp': ['2024-01-01', '2024-01-02', '2024-01-03']
+            })
+
+            # Infer schema for logistics tracking
+            schema = df.ai.infer_schema(
+                purpose="Extract shipping status and location data for logistics tracking"
+            )
+
+            # Apply the schema to extract structured data
+            extracted_df = df.ai.task(schema.task)
+            ```
+
+        Note:
+            The DataFrame rows are internally converted to JSON format before
+            analysis. The inferred schema is flat (no nested structures) and
+            uses only primitive types to ensure compatibility with pandas and
+            Spark operations.
+        """
+        return _df_rows_to_json_series(self._obj).ai.infer_schema(
+            purpose=purpose,
+            max_examples=max_examples,
+        )
+
     def extract(self, column: str) -> pd.DataFrame:
         """Flatten one column of Pydantic models/dicts into top‑level columns.
 
@@ -790,6 +987,100 @@ class OpenAIVecDataFrameAccessor:
 
         return df
 
+    def auto_extract(
+        self,
+        purpose: str,
+        max_examples: int = 100,
+        batch_size: int | None = None,
+        show_progress: bool = False,
+        **api_kwargs,
+    ) -> pd.DataFrame:
+        """Automatically infer schema and add extracted fields to the DataFrame.
+
+        This convenience method combines schema inference and data extraction to
+        automatically add new columns to the existing DataFrame. It analyzes a
+        sample of the DataFrame rows to infer an appropriate schema based on the
+        stated purpose, then extracts structured data and joins it with the
+        original DataFrame.
+
+        Args:
+            purpose (str): Plain language description of what information to extract
+                and how it will be used (e.g., "Extract customer sentiment metrics",
+                "Parse product attributes for analytics"). This guides both schema
+                inference and field selection.
+            max_examples (int): Maximum number of rows to use for schema inference.
+                A larger sample may produce more accurate schemas but increases
+                inference time. Defaults to 100.
+            batch_size (int | None): Number of requests to process in parallel during
+                extraction. Defaults to None (automatic optimization). Set to a specific
+                value to control API usage and performance.
+            show_progress (bool): Whether to display a progress bar during extraction.
+                Useful for large datasets. Defaults to False.
+            **api_kwargs: Additional OpenAI API parameters (e.g., `temperature`, `top_p`,
+                `frequency_penalty`, `presence_penalty`, `seed`) forwarded to the task execution.
+
+        Returns:
+            pd.DataFrame: The original DataFrame with new columns added from the
+                inferred structured data. Each inferred field becomes a new column.
+                The original columns and index are preserved.
+
+        Example:
+            ```python
+            # Add sentiment and issue type to support tickets
+            df = pd.DataFrame({
+                'ticket_id': [1, 2, 3],
+                'description': [
+                    "Can't login, password reset not working",
+                    "Billing error, charged twice last month",
+                    "Great service, issue resolved quickly!"
+                ],
+                'date': ['2024-01-01', '2024-01-02', '2024-01-03']
+            })
+
+            # Add inferred fields to existing DataFrame
+            enriched_df = df.ai.auto_extract(
+                purpose="Extract issue type and sentiment for support dashboard",
+                show_progress=True
+            )
+            # Result: Original df with new columns like 'issue_type', 'sentiment', etc.
+
+            # Add product specifications to inventory data
+            inventory = pd.DataFrame({
+                'sku': ['A001', 'B002', 'C003'],
+                'description': [
+                    "Laptop 16GB RAM, 512GB SSD, Intel i7",
+                    "Phone 128GB, 5G, dual camera",
+                    "Tablet 10-inch, WiFi only, 64GB"
+                ]
+            })
+
+            enriched_inventory = inventory.ai.auto_extract(
+                purpose="Extract technical specifications for inventory system"
+            )
+            ```
+
+        Note:
+            This method is ideal for enriching existing DataFrames with additional
+            structured fields extracted from text columns. The schema is inferred
+            from the entire DataFrame content (converted to JSON format). For
+            production use cases with stable schemas, consider using `infer_schema()`
+            once and reusing the schema with `task()`.
+        """
+        # Infer schema from DataFrame rows
+        schema = self._obj.ai.infer_schema(purpose=purpose, max_examples=max_examples)
+
+        # Extract structured data using the inferred schema
+        inferred_series = self._obj.ai.task(
+            task=schema.task,
+            batch_size=batch_size,
+            show_progress=show_progress,
+            **api_kwargs,
+        )
+
+        return self._obj.assign(
+            inferred=inferred_series,
+        ).ai.extract("inferred")
+
     def similarity(self, col1: str, col2: str) -> pd.Series:
         """Compute cosine similarity between two columns containing embedding vectors.
 
@@ -1165,6 +1456,96 @@ class AsyncOpenAIVecSeriesAccessor:
             **api_kwargs,
         )
 
+    async def auto_extract(
+        self,
+        purpose: str,
+        max_examples: int = 100,
+        batch_size: int | None = None,
+        max_concurrency: int = 8,
+        show_progress: bool = False,
+        **api_kwargs,
+    ) -> pd.DataFrame:
+        """Automatically infer schema and extract structured data in one step (asynchronously).
+
+        This convenience method combines schema inference and data extraction into
+        a single operation. It first analyzes a sample of the Series to infer an
+        appropriate schema based on the stated purpose, then immediately applies
+        that schema to extract structured data from all values in the Series.
+
+        Args:
+            purpose (str): Plain language description of what information to extract
+                and how it will be used (e.g., "Extract product features for search",
+                "Parse customer feedback for sentiment analysis"). This guides both
+                schema inference and field selection.
+            max_examples (int): Maximum number of examples to use for schema inference.
+                A larger sample may produce more accurate schemas but increases
+                inference time. Defaults to 100.
+            batch_size (int | None): Number of requests to process in parallel during
+                extraction. Defaults to None (automatic optimization). Set to a specific
+                value to control API usage and performance.
+            max_concurrency (int): Maximum number of concurrent requests during
+                extraction. Defaults to 8.
+            show_progress (bool): Whether to display a progress bar during extraction.
+                Useful for large datasets. Defaults to False.
+            **api_kwargs: Additional OpenAI API parameters (e.g., `temperature`, `top_p`,
+                `frequency_penalty`, `presence_penalty`, `seed`) forwarded to the task execution.
+
+        Returns:
+            pd.DataFrame: A DataFrame with extracted structured data. Each inferred
+                field becomes a column, with the same index as the original Series.
+                Column names and types are determined by the inferred schema.
+
+        Example:
+            ```python
+            # Extract structured data from product reviews
+            reviews = pd.Series([
+                "Great laptop! 16GB RAM, fast SSD, battery lasts 10 hours",
+                "Decent phone. 128GB storage, camera is okay, screen is bright",
+                "Gaming desktop with RTX 4090, 32GB RAM, runs everything smoothly"
+            ])
+
+            # One-step extraction (must be awaited)
+            extracted = await reviews.aio.auto_extract(
+                purpose="Extract product specifications and performance metrics",
+                max_concurrency=4,
+                show_progress=True
+            )
+            # Result: DataFrame with columns like 'ram', 'storage', 'battery_life', etc.
+
+            # Extract sentiment and issues from support tickets
+            tickets = pd.Series([
+                "Account locked, can't reset password, very frustrated",
+                "Billing error, charged twice for subscription",
+                "Great support! Issue resolved quickly"
+            ])
+
+            features = await tickets.aio.auto_extract(
+                purpose="Extract issue type and customer sentiment for support analytics",
+                batch_size=32
+            )
+            ```
+
+        Note:
+            This is an asynchronous method and must be awaited. This method is ideal
+            for exploratory data analysis when you don't have a predefined schema.
+            For production use cases with stable schemas, consider using the synchronous
+            `infer_schema()` once and reusing the schema with `task()`. The inferred
+            schema is not returned, so if you need to inspect or save it, use
+            `infer_schema()` and `task()` separately.
+        """
+        # Use synchronous infer_schema since it's not async
+        schema = self._obj.ai.infer_schema(purpose=purpose, max_examples=max_examples)
+
+        inferred_series = await self._obj.aio.task(
+            task=schema.task,
+            batch_size=batch_size,
+            max_concurrency=max_concurrency,
+            show_progress=show_progress,
+            **api_kwargs,
+        )
+
+        return pd.DataFrame({"inferred": inferred_series}).ai.extract("inferred")
+
 
 @pd.api.extensions.register_dataframe_accessor("aio")
 class AsyncOpenAIVecDataFrameAccessor:
@@ -1572,3 +1953,103 @@ class AsyncOpenAIVecDataFrameAccessor:
                 df.at[actual_index, target_column_name] = result.output
 
         return df
+
+    async def auto_extract(
+        self,
+        purpose: str,
+        max_examples: int = 100,
+        batch_size: int | None = None,
+        max_concurrency: int = 8,
+        show_progress: bool = False,
+        **api_kwargs,
+    ) -> pd.DataFrame:
+        """Automatically infer schema and add extracted fields to the DataFrame (asynchronously).
+
+        This convenience method combines schema inference and data extraction to
+        automatically add new columns to the existing DataFrame. It analyzes a
+        sample of the DataFrame rows to infer an appropriate schema based on the
+        stated purpose, then extracts structured data and joins it with the
+        original DataFrame.
+
+        Args:
+            purpose (str): Plain language description of what information to extract
+                and how it will be used (e.g., "Extract customer sentiment metrics",
+                "Parse product attributes for analytics"). This guides both schema
+                inference and field selection.
+            max_examples (int): Maximum number of rows to use for schema inference.
+                A larger sample may produce more accurate schemas but increases
+                inference time. Defaults to 100.
+            batch_size (int | None): Number of requests to process in parallel during
+                extraction. Defaults to None (automatic optimization). Set to a specific
+                value to control API usage and performance.
+            max_concurrency (int): Maximum number of concurrent requests during
+                extraction. Defaults to 8.
+            show_progress (bool): Whether to display a progress bar during extraction.
+                Useful for large datasets. Defaults to False.
+            **api_kwargs: Additional OpenAI API parameters (e.g., `temperature`, `top_p`,
+                `frequency_penalty`, `presence_penalty`, `seed`) forwarded to the task execution.
+
+        Returns:
+            pd.DataFrame: The original DataFrame with new columns added from the
+                inferred structured data. Each inferred field becomes a new column.
+                The original columns and index are preserved.
+
+        Example:
+            ```python
+            # Add sentiment and issue type to support tickets
+            df = pd.DataFrame({
+                'ticket_id': [1, 2, 3],
+                'description': [
+                    "Can't login, password reset not working",
+                    "Billing error, charged twice last month",
+                    "Great service, issue resolved quickly!"
+                ],
+                'date': ['2024-01-01', '2024-01-02', '2024-01-03']
+            })
+
+            # Add inferred fields to existing DataFrame (must be awaited)
+            enriched_df = await df.aio.auto_extract(
+                purpose="Extract issue type and sentiment for support dashboard",
+                max_concurrency=4,
+                show_progress=True
+            )
+            # Result: Original df with new columns like 'issue_type', 'sentiment', etc.
+
+            # Add product specifications to inventory data
+            inventory = pd.DataFrame({
+                'sku': ['A001', 'B002', 'C003'],
+                'description': [
+                    "Laptop 16GB RAM, 512GB SSD, Intel i7",
+                    "Phone 128GB, 5G, dual camera",
+                    "Tablet 10-inch, WiFi only, 64GB"
+                ]
+            })
+
+            enriched_inventory = await inventory.aio.auto_extract(
+                purpose="Extract technical specifications for inventory system",
+                batch_size=32
+            )
+            ```
+
+        Note:
+            This is an asynchronous method and must be awaited. This method is ideal
+            for enriching existing DataFrames with additional structured fields
+            extracted from text columns. The schema is inferred synchronously from
+            the DataFrame content. For production use cases with stable schemas,
+            consider using `infer_schema()` once and reusing the schema with `task()`.
+        """
+        # Infer schema from DataFrame rows (synchronous)
+        schema = self._obj.ai.infer_schema(purpose=purpose, max_examples=max_examples)
+
+        # Extract structured data using the inferred schema (asynchronous)
+        inferred_series = await self._obj.aio.task(
+            task=schema.task,
+            batch_size=batch_size,
+            max_concurrency=max_concurrency,
+            show_progress=show_progress,
+            **api_kwargs,
+        )
+
+        return self._obj.assign(
+            inferred=inferred_series,
+        ).ai.extract("inferred")