undatum 1.0.17__py2.py3-none-any.whl

undatum/ai/providers.py

@@ -0,0 +1,1002 @@
+"""AI service provider implementations."""
+import json
+import os
+import re
+import time
+from typing import Dict, Optional, Callable
+
+import requests
+
+from .base import AIService, AIConfigurationError, AIAPIError
+from .schemas import FIELD_INFO_SCHEMA, DATASET_DESCRIPTION_SCHEMA
+
+
+def retry_with_backoff(func: Callable, max_retries: int = 3, initial_delay: float = 1.0,
+                      backoff_factor: float = 2.0, retry_statuses: tuple = (429, 500, 502, 503, 504)):
+    """Retry a function with exponential backoff for rate limiting and server errors.
+
+    Args:
+        func: Function to retry (should raise requests.exceptions.RequestException)
+        max_retries: Maximum number of retry attempts
+        initial_delay: Initial delay in seconds before first retry
+        backoff_factor: Factor to multiply delay by for each retry
+        retry_statuses: HTTP status codes that should trigger a retry
+
+    Returns:
+        Result of the function call
+
+    Raises:
+        AIAPIError: If all retries are exhausted
+    """
+    delay = initial_delay
+    last_exception = None
+
+    for attempt in range(max_retries + 1):
+        try:
+            return func()
+        except requests.exceptions.RequestException as e:
+            last_exception = e
+            status_code = getattr(e.response, 'status_code', None) if hasattr(e, 'response') else None
+
+            # Only retry on specific status codes or if status_code is None (network error)
+            if status_code not in retry_statuses and status_code is not None:
+                raise
+
+            # Don't retry on last attempt
+            if attempt >= max_retries:
+                break
+
+            # Extract retry-after header if present (for 429 errors)
+            retry_after = None
+            if hasattr(e, 'response') and e.response is not None:
+                retry_after = e.response.headers.get('Retry-After')
+                if retry_after:
+                    try:
+                        delay = float(retry_after)
+                    except ValueError:
+                        pass
+
+            # Wait before retrying
+            time.sleep(delay)
+            delay *= backoff_factor
+
+    # All retries exhausted, raise the last exception
+    status_code = getattr(last_exception.response, 'status_code', None) if hasattr(last_exception, 'response') else None
+    error_msg = f"API request failed after {max_retries + 1} attempts: {str(last_exception)}"
+
+    if status_code == 429:
+        error_msg += "\nRate limit exceeded. Please wait a moment and try again, or check your API usage limits."
+    elif status_code in (500, 502, 503, 504):
+        error_msg += "\nServer error. The API may be temporarily unavailable. Please try again later."
+
+    raise AIAPIError(error_msg,
+                    status_code=status_code,
+                    response=getattr(last_exception.response, 'text', None) if hasattr(last_exception, 'response') else None)
+
+
+class OpenAIProvider(AIService):
+    """OpenAI API provider."""
+
+    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None,
+                 model: Optional[str] = None, timeout: int = 30):
+        """Initialize OpenAI provider.
+
+        Args:
+            api_key: OpenAI API key (defaults to OPENAI_API_KEY env var)
+            base_url: Base URL (defaults to https://api.openai.com/v1)
+            model: Model name (defaults to gpt-4o-mini)
+            timeout: Request timeout in seconds
+        """
+        api_key = api_key or os.getenv('OPENAI_API_KEY')
+        base_url = base_url or 'https://api.openai.com/v1'
+        model = model or 'gpt-4o-mini'
+        super().__init__(api_key, base_url, model, timeout)
+        self._validate_config()
+
+    def get_fields_info(self, fields: list[str], language: str = 'English') -> Dict[str, str]:
+        """Get field descriptions using OpenAI API."""
+        fields_str = ', '.join(fields)
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide clear, concise descriptions of data fields in {language}. Always respond with valid JSON."
+                },
+                {
+                    "role": "user",
+                    "content": f"Please describe these data fields in {language}: {fields_str}. Provide a description for each field explaining what it represents. Return your response as a JSON object with a 'fields' array containing objects with 'name' and 'description' keys."
+                }
+            ],
+            "response_format": {
+                "type": "json_object"
+            },
+            "temperature": 0.3
+        }
+
+        def _make_request():
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            return response
+
+        try:
+            response = retry_with_backoff(_make_request)
+            data = response.json()
+
+            content = data["choices"][0]["message"]["content"]
+            result = json.loads(content)
+
+            # Validate and convert to expected format
+            if "fields" not in result:
+                raise AIAPIError("Invalid response format: missing 'fields' key")
+
+            field_dict = {}
+            for field_info in result["fields"]:
+                if "name" not in field_info or "description" not in field_info:
+                    continue
+                field_dict[field_info["name"]] = field_info["description"]
+
+            # Ensure all requested fields are in the result
+            for field in fields:
+                if field not in field_dict:
+                    field_dict[field] = f"Field: {field}"
+
+            return field_dict
+
+        except AIAPIError:
+            raise
+        except requests.exceptions.RequestException as e:
+            raise AIAPIError(f"OpenAI API request failed: {str(e)}",
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError as e:
+            raise AIAPIError(f"Failed to parse OpenAI response: {str(e)}")
+
+    def get_description(self, data: str, language: str = 'English') -> str:
+        """Get dataset description using OpenAI API."""
+        # Truncate data if too large (OpenAI has token limits)
+        # Use a more conservative limit to account for prompt overhead
+        MAX_DATA_LENGTH = 3000
+        if len(data) > MAX_DATA_LENGTH:
+            data = data[:MAX_DATA_LENGTH] + "\n... (truncated)"
+
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+
+        user_content = f"""I have the following CSV data sample:
+{data}
+Please provide a short description of this dataset in {language}. Consider this as a sample of a larger dataset. Don't generate code or data examples.
+Return your response as a JSON object with a "description" key."""
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide concise dataset descriptions in {language}. Always respond with valid JSON."
+                },
+                {
+                    "role": "user",
+                    "content": user_content
+                }
+            ],
+            "response_format": {
+                "type": "json_object"
+            },
+            "temperature": 0.3
+        }
+
+        def _make_request():
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            return response
+
+        try:
+            response = retry_with_backoff(_make_request)
+            response_data = response.json()
+
+            content = response_data["choices"][0]["message"]["content"]
+            result = json.loads(content)
+
+            if "description" in result:
+                return result["description"]
+            else:
+                # Fallback: return the content as-is if structure is different
+                return content
+
+        except AIAPIError:
+            raise
+        except requests.exceptions.RequestException as e:
+            error_msg = f"OpenAI API request failed: {str(e)}"
+            if hasattr(e, 'response') and e.response is not None:
+                try:
+                    error_detail = e.response.json()
+                    if 'error' in error_detail:
+                        error_info = error_detail['error']
+                        if 'message' in error_info:
+                            error_msg += f"\nError details: {error_info['message']}"
+                        if 'code' in error_info:
+                            error_msg += f"\nError code: {error_info['code']}"
+                except (ValueError, KeyError):
+                    # If we can't parse the error response, include the raw text
+                    error_text = getattr(e.response, 'text', None)
+                    if error_text:
+                        error_msg += f"\nResponse: {error_text[:500]}"
+            raise AIAPIError(error_msg,
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError:
+            # If JSON parsing fails, try to extract description from text
+            try:
+                content = response_data["choices"][0]["message"]["content"]
+                return content
+            except (KeyError, IndexError):
+                raise AIAPIError("Failed to extract description from OpenAI response")
+
+
+class OpenRouterProvider(AIService):
+    """OpenRouter API provider (OpenAI-compatible)."""
+
+    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None,
+                 model: Optional[str] = None, timeout: int = 30):
+        """Initialize OpenRouter provider.
+
+        Args:
+            api_key: OpenRouter API key (defaults to OPENROUTER_API_KEY env var)
+            base_url: Base URL (defaults to https://openrouter.ai/api/v1)
+            model: Model name (defaults to openai/gpt-4o-mini)
+            timeout: Request timeout in seconds
+        """
+        api_key = api_key or os.getenv('OPENROUTER_API_KEY')
+        base_url = base_url or 'https://openrouter.ai/api/v1'
+        model = model or 'openai/gpt-4o-mini'
+        super().__init__(api_key, base_url, model, timeout)
+        self._validate_config()
+
+    def get_fields_info(self, fields: list[str], language: str = 'English') -> Dict[str, str]:
+        """Get field descriptions using OpenRouter API."""
+        fields_str = ', '.join(fields)
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://github.com/datacoon/undatum",
+            "X-Title": "Undatum Data Analysis"
+        }
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide clear, concise descriptions of data fields in {language}. Always respond with valid JSON."
+                },
+                {
+                    "role": "user",
+                    "content": f"Please describe these data fields in {language}: {fields_str}. Provide a description for each field explaining what it represents. Return your response as a JSON object with a 'fields' array containing objects with 'name' and 'description' keys."
+                }
+            ],
+            "response_format": {
+                "type": "json_object"
+            },
+            "temperature": 0.3
+        }
+
+        def _make_request():
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            return response
+
+        try:
+            response = retry_with_backoff(_make_request)
+            data = response.json()
+
+            content = data["choices"][0]["message"]["content"]
+            result = json.loads(content)
+
+            field_dict = {}
+            if "fields" in result:
+                for field_info in result["fields"]:
+                    if "name" in field_info and "description" in field_info:
+                        field_dict[field_info["name"]] = field_info["description"]
+
+            # Ensure all requested fields are in the result
+            for field in fields:
+                if field not in field_dict:
+                    field_dict[field] = f"Field: {field}"
+
+            return field_dict
+
+        except AIAPIError:
+            raise
+        except requests.exceptions.RequestException as e:
+            raise AIAPIError(f"OpenRouter API request failed: {str(e)}",
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError as e:
+            raise AIAPIError(f"Failed to parse OpenRouter response: {str(e)}")
+
+    def get_description(self, data: str, language: str = 'English') -> str:
+        """Get dataset description using OpenRouter API."""
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+            "HTTP-Referer": "https://github.com/datacoon/undatum",
+            "X-Title": "Undatum Data Analysis"
+        }
+
+        user_content = f"""I have the following CSV data sample:
+{data}
+Please provide a short description of this dataset in {language}. Consider this as a sample of a larger dataset. Don't generate code or data examples.
+Return your response as a JSON object with a "description" key."""
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide concise dataset descriptions in {language}. Always respond with valid JSON."
+                },
+                {
+                    "role": "user",
+                    "content": user_content
+                }
+            ],
+            "response_format": {
+                "type": "json_object"
+            },
+            "temperature": 0.3
+        }
+
+        # Truncate data if too large
+        MAX_DATA_LENGTH = 5000
+        if len(data) > MAX_DATA_LENGTH:
+            data = data[:MAX_DATA_LENGTH] + "\n... (truncated)"
+            payload["messages"][1]["content"] = f"""I have the following CSV data sample:
+{data}
+Please provide a short description of this dataset in {language}. Consider this as a sample of a larger dataset. Don't generate code or data examples. Return JSON with a 'description' key."""
+
+        def _make_request():
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            return response
+
+        try:
+            response = retry_with_backoff(_make_request)
+            response_data = response.json()
+
+            content = response_data["choices"][0]["message"]["content"]
+            result = json.loads(content)
+
+            if "description" in result:
+                return result["description"]
+            else:
+                return content
+
+        except AIAPIError:
+            raise
+        except requests.exceptions.RequestException as e:
+            raise AIAPIError(f"OpenRouter API request failed: {str(e)}",
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError:
+            try:
+                content = response_data["choices"][0]["message"]["content"]
+                return content
+            except (KeyError, IndexError):
+                raise AIAPIError("Failed to extract description from OpenRouter response")
+
+
+class OllamaProvider(AIService):
+    """Ollama local API provider."""
+
+    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None,
+                 model: Optional[str] = None, timeout: int = 30):
+        """Initialize Ollama provider.
+
+        Args:
+            api_key: Not used for Ollama (kept for interface compatibility)
+            base_url: Base URL (defaults to http://localhost:11434)
+            model: Model name (defaults to llama3.2)
+            timeout: Request timeout in seconds
+        """
+        base_url = base_url or os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434')
+        model = model or 'llama3.2'
+        super().__init__(api_key, base_url, model, timeout)
+        if not self.model:
+            raise AIConfigurationError("Model is required for OllamaProvider")
+
+    def _validate_config(self) -> None:
+        """Ollama doesn't require API key."""
+        if not self.model:
+            raise AIConfigurationError("Model is required for OllamaProvider")
+
+    def get_fields_info(self, fields: list[str], language: str = 'English') -> Dict[str, str]:
+        """Get field descriptions using Ollama API."""
+        fields_str = ', '.join(fields)
+        url = f"{self.base_url}/api/chat"
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide clear, concise descriptions of data fields in {language}. Always respond with valid JSON only."
+                },
+                {
+                    "role": "user",
+                    "content": f"""Please describe these data fields in {language}: {fields_str}.
+Return a JSON object with a "fields" array. Each item should have "name" and "description" keys.
+Example format: {{"fields": [{{"name": "field1", "description": "..."}}, {{"name": "field2", "description": "..."}}]}}"""
+                }
+            ],
+            "format": "json",
+            "stream": False,  # Explicitly disable streaming
+            "options": {
+                "temperature": 0.3
+            }
+        }
+
+        try:
+            response = requests.post(url, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            
+            # Handle potential streaming response or malformed JSON
+            # Ollama may return multiple JSON objects even with stream=False
+            response_text = response.text.strip()
+            data = None
+            
+            try:
+                data = response.json()
+            except (ValueError, json.JSONDecodeError) as e:
+                # If response.json() fails, it might be a streaming response with multiple JSON objects
+                # Try to parse the last complete JSON object from the response text
+                if response_text:
+                    # Split by newlines and try to parse each line as JSON
+                    # Ollama streaming format has one JSON object per line
+                    lines = response_text.strip().split('\n')
+                    for line in reversed(lines):  # Start from the last line
+                        line = line.strip()
+                        if not line:
+                            continue
+                        try:
+                            # Try to parse as JSON
+                            parsed = json.loads(line)
+                            # Check if it's a valid Ollama response structure
+                            if isinstance(parsed, dict) and ('message' in parsed or 'content' in parsed or 'response' in parsed):
+                                data = parsed
+                                break
+                        except json.JSONDecodeError:
+                            continue
+                    
+                    # If we still don't have data, try to extract JSON object with regex
+                    if data is None:
+                        json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
+                        if json_match:
+                            try:
+                                data = json.loads(json_match.group(0))
+                            except json.JSONDecodeError:
+                                pass
+                    
+                    if data is None:
+                        raise AIAPIError(f"Failed to parse Ollama response: {str(e)}. Response: {response_text[:500]}")
+                else:
+                    raise AIAPIError(f"Empty response from Ollama: {str(e)}")
+
+            content = data.get("message", {}).get("content", "")
+            
+            # Try to extract JSON from content if it contains markdown code blocks
+            json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
+            if json_match:
+                content = json_match.group(1)
+            else:
+                # Try to find JSON object in the text
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    content = json_match.group(0)
+            
+            result = json.loads(content)
+
+            field_dict = {}
+            if "fields" in result:
+                for field_info in result["fields"]:
+                    if "name" in field_info and "description" in field_info:
+                        field_dict[field_info["name"]] = field_info["description"]
+
+            # Ensure all requested fields are in the result
+            for field in fields:
+                if field not in field_dict:
+                    field_dict[field] = f"Field: {field}"
+
+            return field_dict
+
+        except requests.exceptions.RequestException as e:
+            raise AIAPIError(f"Ollama API request failed: {str(e)}",
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError as e:
+            raise AIAPIError(f"Failed to parse Ollama response: {str(e)}. Content: {content[:200] if 'content' in locals() else 'N/A'}")
+
+    def get_description(self, data: str, language: str = 'English') -> str:
+        """Get dataset description using Ollama API."""
+        url = f"{self.base_url}/api/chat"
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide concise dataset descriptions in {language}. Always respond with valid JSON only."
+                },
+                {
+                    "role": "user",
+                    "content": f"""I have the following CSV data sample:
+{data}
+Please provide a short description of this dataset in {language}. Consider this as a sample of a larger dataset. Don't generate code or data examples.
+Return JSON with format: {{"description": "..."}}"""
+                }
+            ],
+            "format": "json",
+            "stream": False,  # Explicitly disable streaming
+            "options": {
+                "temperature": 0.3
+            }
+        }
+
+        try:
+            response = requests.post(url, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            
+            # Handle potential streaming response or malformed JSON
+            # Ollama may return multiple JSON objects even with stream=False
+            response_text = response.text.strip()
+            data = None
+            
+            try:
+                data = response.json()
+            except (ValueError, json.JSONDecodeError) as e:
+                # If response.json() fails, it might be a streaming response with multiple JSON objects
+                # Try to parse the last complete JSON object from the response text
+                if response_text:
+                    # Split by newlines and try to parse each line as JSON
+                    # Ollama streaming format has one JSON object per line
+                    lines = response_text.strip().split('\n')
+                    for line in reversed(lines):  # Start from the last line
+                        line = line.strip()
+                        if not line:
+                            continue
+                        try:
+                            # Try to parse as JSON
+                            parsed = json.loads(line)
+                            # Check if it's a valid Ollama response structure
+                            if isinstance(parsed, dict) and ('message' in parsed or 'content' in parsed or 'response' in parsed):
+                                data = parsed
+                                break
+                        except json.JSONDecodeError:
+                            continue
+                    
+                    # If we still don't have data, try to extract JSON object with regex
+                    if data is None:
+                        json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
+                        if json_match:
+                            try:
+                                data = json.loads(json_match.group(0))
+                            except json.JSONDecodeError:
+                                pass
+                    
+                    if data is None:
+                        raise AIAPIError(f"Failed to parse Ollama response: {str(e)}. Response: {response_text[:500]}")
+                else:
+                    raise AIAPIError(f"Empty response from Ollama: {str(e)}")
+
+            content = data.get("message", {}).get("content", "")
+            
+            # Try to extract JSON from content if it contains markdown code blocks
+            json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
+            if json_match:
+                content = json_match.group(1)
+            else:
+                # Try to find JSON object in the text
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    content = json_match.group(0)
+            
+            # Try to parse as JSON
+            try:
+                result = json.loads(content)
+                if "description" in result:
+                    return result["description"]
+                else:
+                    # If JSON is valid but doesn't have description, return content as-is
+                    return content
+            except json.JSONDecodeError as json_err:
+                # If JSON parsing fails, try to return the raw content
+                # This handles cases where the model returns plain text instead of JSON
+                if content:
+                    return content
+                raise AIAPIError(f"Failed to parse Ollama JSON response: {str(json_err)}. Content: {content[:200]}")
+
+        except requests.exceptions.RequestException as e:
+            raise AIAPIError(f"Ollama API request failed: {str(e)}",
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError as e:
+            # This should not happen now, but keep as fallback
+            try:
+                content = data.get("message", {}).get("content", "")
+                if content:
+                    return content
+            except (KeyError, IndexError, NameError):
+                pass
+            raise AIAPIError(f"Failed to extract description from Ollama response: {str(e)}")
+
+
+class LMStudioProvider(AIService):
+    """LM Studio local API provider (OpenAI-compatible)."""
+
+    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None,
+                 model: Optional[str] = None, timeout: int = 30):
+        """Initialize LM Studio provider.
+
+        Args:
+            api_key: Not used for LM Studio (kept for interface compatibility, can be "lm-studio")
+            base_url: Base URL (defaults to http://localhost:1234/v1)
+            model: Model name (REQUIRED - must match a model loaded in LM Studio)
+            timeout: Request timeout in seconds
+        """
+        base_url = base_url or os.getenv('LMSTUDIO_BASE_URL', 'http://localhost:1234/v1')
+        super().__init__(api_key, base_url, model, timeout)
+        self._validate_config()
+
+    def _validate_config(self) -> None:
+        """Validate LM Studio configuration."""
+        if not self.model:
+            raise AIConfigurationError(
+                "Model name is required for LM Studio. "
+                "Please specify a model name that matches a model loaded in LM Studio. "
+                "Example: --ai-model 'your-model-name'"
+            )
+
+    def _get_available_models(self) -> list[str]:
+        """Get list of available models from LM Studio."""
+        try:
+            models_url = f"{self.base_url}/models"
+            response = requests.get(models_url, timeout=self.timeout)
+            if response.status_code == 200:
+                data = response.json()
+                if isinstance(data, dict) and 'data' in data:
+                    return [model.get('id', '') for model in data['data'] if 'id' in model]
+        except Exception:
+            pass
+        return []
+
+    def get_fields_info(self, fields: list[str], language: str = 'English') -> Dict[str, str]:
+        """Get field descriptions using LM Studio API."""
+        fields_str = ', '.join(fields)
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Content-Type": "application/json"
+        }
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide clear, concise descriptions of data fields in {language}. Always respond with valid JSON only, no markdown, no code blocks."
+                },
+                {
+                    "role": "user",
+                    "content": f"""Please describe these data fields in {language}: {fields_str}.
+Return ONLY a JSON object with a "fields" array. Each item must have "name" and "description" keys.
+Format: {{"fields": [{{"name": "field1", "description": "..."}}, {{"name": "field2", "description": "..."}}]}}
+Return only the JSON, nothing else."""
+                }
+            ],
+            "temperature": 0.3
+        }
+
+        # Try with json_object format first (some models support it)
+        # If that fails, fall back to text parsing
+        try:
+            payload["response_format"] = {"type": "json_object"}
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+        except requests.exceptions.RequestException:
+            # Remove response_format if not supported
+            payload.pop("response_format", None)
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+
+        try:
+            data = response.json()
+            content = data["choices"][0]["message"]["content"].strip()
+
+            # Try to extract JSON from markdown code blocks if present
+            json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
+            if json_match:
+                content = json_match.group(1)
+            else:
+                # Try to find JSON object in the text
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    content = json_match.group(0)
+
+            result = json.loads(content)
+
+            field_dict = {}
+            if "fields" in result:
+                for field_info in result["fields"]:
+                    if "name" in field_info and "description" in field_info:
+                        field_dict[field_info["name"]] = field_info["description"]
+
+            # Ensure all requested fields are in the result
+            for field in fields:
+                if field not in field_dict:
+                    field_dict[field] = f"Field: {field}"
+
+            return field_dict
+
+        except requests.exceptions.RequestException as e:
+            available_models = self._get_available_models()
+            error_msg = f"LM Studio API request failed: {str(e)}"
+            if available_models:
+                error_msg += f"\nAvailable models: {', '.join(available_models)}"
+            elif e.response and e.response.status_code == 404:
+                error_msg += "\nMake sure LM Studio server is running and a model is loaded."
+            raise AIAPIError(error_msg,
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError as e:
+            raise AIAPIError(f"Failed to parse LM Studio response: {str(e)}. Response content: {content[:200]}")
+
+    def get_description(self, data: str, language: str = 'English') -> str:
+        """Get dataset description using LM Studio API."""
+        # Truncate data if too large
+        MAX_DATA_LENGTH = 5000
+        if len(data) > MAX_DATA_LENGTH:
+            data = data[:MAX_DATA_LENGTH] + "\n... (truncated)"
+
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Content-Type": "application/json"
+        }
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide concise dataset descriptions in {language}. Always respond with valid JSON only, no markdown, no code blocks."
+                },
+                {
+                    "role": "user",
+                    "content": f"""I have the following CSV data sample:
+{data}
+Please provide a short description of this dataset in {language}. Consider this as a sample of a larger dataset. Don't generate code or data examples.
+Return ONLY a JSON object with format: {{"description": "..."}}
+Return only the JSON, nothing else."""
+                }
+            ],
+            "temperature": 0.3
+        }
+
+        # Try with json_object format first (some models support it)
+        # If that fails, fall back to text parsing
+        try:
+            payload["response_format"] = {"type": "json_object"}
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+        except requests.exceptions.RequestException:
+            # Remove response_format if not supported
+            payload.pop("response_format", None)
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+
+        try:
+            response_data = response.json()
+            content = response_data["choices"][0]["message"]["content"].strip()
+
+            # Try to extract JSON from markdown code blocks if present
+            json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
+            if json_match:
+                content = json_match.group(1)
+            else:
+                # Try to find JSON object in the text
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    content = json_match.group(0)
+
+            # Try to parse as JSON first
+            try:
+                result = json.loads(content)
+                if "description" in result:
+                    return result["description"]
+            except json.JSONDecodeError:
+                pass
+
+            # If JSON parsing fails, return the content as-is (might be plain text description)
+            return content
+
+        except requests.exceptions.RequestException as e:
+            available_models = self._get_available_models()
+            error_msg = f"LM Studio API request failed: {str(e)}"
+            if available_models:
+                error_msg += f"\nAvailable models: {', '.join(available_models)}"
+            elif e.response and e.response.status_code == 404:
+                error_msg += "\nMake sure LM Studio server is running and a model is loaded."
+            raise AIAPIError(error_msg,
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except (KeyError, IndexError) as e:
+            raise AIAPIError(f"Failed to extract description from LM Studio response: {str(e)}")
+
+
+class PerplexityProvider(AIService):
+    """Perplexity API provider with structured output."""
+
+    def __init__(self, api_key: Optional[str] = None, base_url: Optional[str] = None,
+                 model: Optional[str] = None, timeout: int = 30):
+        """Initialize Perplexity provider.
+
+        Args:
+            api_key: Perplexity API key (defaults to PERPLEXITY_API_KEY env var)
+            base_url: Base URL (defaults to https://api.perplexity.ai)
+            model: Model name (defaults to sonar)
+            timeout: Request timeout in seconds
+        """
+        # Always prioritize PERPLEXITY_API_KEY environment variable
+        perplexity_key = os.getenv('PERPLEXITY_API_KEY')
+        if perplexity_key:
+            # If PERPLEXITY_API_KEY is set, always use it (ignore passed api_key)
+            api_key = perplexity_key
+        # If PERPLEXITY_API_KEY is not set, use the passed api_key (which may be None)
+
+        base_url = base_url or 'https://api.perplexity.ai'
+        model = model or 'sonar'
+        super().__init__(api_key, base_url, model, timeout)
+        self._validate_config()
+
+    def get_fields_info(self, fields: list[str], language: str = 'English') -> Dict[str, str]:
+        """Get field descriptions using Perplexity API."""
+        fields_str = ', '.join(fields)
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide clear, concise descriptions of data fields in {language}. Always respond with valid JSON only, no markdown, no code blocks."
+                },
+                {
+                    "role": "user",
+                    "content": f"""Please describe these data fields in {language}: {fields_str}.
+Return ONLY a JSON object with a "fields" array. Each item must have "name" and "description" keys.
+Format: {{"fields": [{{"name": "field1", "description": "..."}}, {{"name": "field2", "description": "..."}}]}}
+Return only the JSON, nothing else."""
+                }
+            ],
+            "temperature": 0.3
+        }
+
+        try:
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            data = response.json()
+
+            content = data["choices"][0]["message"]["content"].strip()
+
+            # Try to extract JSON from markdown code blocks if present
+            json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
+            if json_match:
+                content = json_match.group(1)
+            else:
+                # Try to find JSON object in the text
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    content = json_match.group(0)
+
+            result = json.loads(content)
+
+            field_dict = {}
+            if "fields" in result:
+                for field_info in result["fields"]:
+                    if "name" in field_info and "description" in field_info:
+                        field_dict[field_info["name"]] = field_info["description"]
+
+            # Ensure all requested fields are in the result
+            for field in fields:
+                if field not in field_dict:
+                    field_dict[field] = f"Field: {field}"
+
+            return field_dict
+
+        except requests.exceptions.RequestException as e:
+            raise AIAPIError(f"Perplexity API request failed: {str(e)}",
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except json.JSONDecodeError as e:
+            raise AIAPIError(f"Failed to parse Perplexity response: {str(e)}. Response content: {content[:200]}")
+
+    def get_description(self, data: str, language: str = 'English') -> str:
+        """Get dataset description using Perplexity API."""
+        # Truncate data if too large (Perplexity has token limits)
+        # Keep first ~5000 characters to ensure we stay within limits
+        MAX_DATA_LENGTH = 5000
+        if len(data) > MAX_DATA_LENGTH:
+            data = data[:MAX_DATA_LENGTH] + "\n... (truncated)"
+
+        url = f"{self.base_url}/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+
+        payload = {
+            "model": self.model,
+            "messages": [
+                {
+                    "role": "system",
+                    "content": f"You are a data documentation assistant. Provide concise dataset descriptions in {language}. Always respond with valid JSON only, no markdown, no code blocks."
+                },
+                {
+                    "role": "user",
+                    "content": f"""I have the following CSV data sample:
+{data}
+Please provide a short description of this dataset in {language}. Consider this as a sample of a larger dataset. Don't generate code or data examples.
+Return ONLY a JSON object with format: {{"description": "..."}}
+Return only the JSON, nothing else."""
+                }
+            ],
+            "temperature": 0.3
+        }
+
+        try:
+            response = requests.post(url, headers=headers, json=payload, timeout=self.timeout)
+            response.raise_for_status()
+            response_data = response.json()
+
+            content = response_data["choices"][0]["message"]["content"].strip()
+
+            # Try to extract JSON from markdown code blocks if present
+            json_match = re.search(r'```(?:json)?\s*(\{.*?\})\s*```', content, re.DOTALL)
+            if json_match:
+                content = json_match.group(1)
+            else:
+                # Try to find JSON object in the text
+                json_match = re.search(r'\{.*\}', content, re.DOTALL)
+                if json_match:
+                    content = json_match.group(0)
+
+            # Try to parse as JSON first
+            try:
+                result = json.loads(content)
+                if "description" in result:
+                    return result["description"]
+            except json.JSONDecodeError:
+                pass
+
+            # If JSON parsing fails, return the content as-is (might be plain text description)
+            return content
+
+        except requests.exceptions.RequestException as e:
+            raise AIAPIError(f"Perplexity API request failed: {str(e)}",
+                           status_code=getattr(e.response, 'status_code', None),
+                           response=getattr(e.response, 'text', None))
+        except (KeyError, IndexError) as e:
+            raise AIAPIError(f"Failed to extract description from Perplexity response: {str(e)}")