sovai 0.2.76__tar.gz → 0.2.77__tar.gz

{sovai-0.2.76 → sovai-0.2.77}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sovai
-Version: 0.2.76
+Version: 0.2.77
 Summary: python for asset management
 License: Proprietary License
 License-File: LICENSE

{sovai-0.2.76 → sovai-0.2.77}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "sovai"
-version = "0.2.76"
+version = "0.2.77"
 description = "python for asset management"
 authors = [
     {name = "Derek Snow", email = "d.snow@outlook.com"}

{sovai-0.2.76 → sovai-0.2.77}/sovai/api_config.py

@@ -14,6 +14,7 @@ class ApiConfig:
     base_url: str = "https://data.sov.ai"
     token_type: str = "Bearer"
     verify_ssl: bool = True
+    version: Optional[str] = None
 
 
 def read_key(envpath: Union[str, Path]) -> Optional[str]:
@@ -45,7 +46,9 @@ def save_key(envpath: Union[str, Path] = ".env") -> None:
         envpath = Path(envpath)
 
     with open(envpath, "w", encoding="utf8") as fh:
-        fh.write(f"API_TOKEN={ApiConfig.token}")
+        fh.write(f"API_TOKEN={ApiConfig.token}\n")
+        if ApiConfig.version:
+            fh.write(f"API_VERSION={ApiConfig.version}\n")
 
 
 def _lookup_token(config: dict[str, Optional[str]]) -> Optional[str]:

{sovai-0.2.76 → sovai-0.2.77}/sovai/extensions/signal_evaluation.py

@@ -1182,7 +1182,10 @@ def stress_plotting(portfolio_returns):
     start_date, end_date = min(all_dates), max(all_dates)
 
     # Add horizontal lines for average returns of selected events
+    # Only process events that exist in event_avg_returns (handles shorter time series)
     for i, event in enumerate(selected_events):
+        if event not in event_avg_returns:
+            continue  # Skip events not present in the data
         avg_return = event_avg_returns[event]
         color = color_map[list(event_returns.keys()).index(event) % len(color_map)]
         fig.add_trace(

{sovai-0.2.76 → sovai-0.2.77}/sovai/get_data.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union, Tuple, List, Dict, Any, Callable, TypeVar
+from typing import Optional, Union, Tuple, List, Dict, Any, Callable, TypeVar, TYPE_CHECKING
 import re
 import time
 import functools
@@ -20,7 +20,6 @@ from requests.exceptions import RequestException, ConnectionError, Timeout, HTTP
 import boto3
 import polars as pl
 import pyarrow.parquet as pq
-import plotly.graph_objects as go
 
 
 # Local imports
@@ -41,16 +40,52 @@ logger = logging.getLogger(__name__)
 T = TypeVar('T')
 DataFrameType = Union[pd.DataFrame, pl.DataFrame, 'CustomDataFrame']
 
+# Version-based column exclusions configuration
+VERSION_COLUMN_EXCLUSIONS = {
+    "p72": {
+        "lobbying": {
+            "always": [
+                    "registrant_name",
+                    "registrant_address",
+                    "registrant_contact_phone",
+                    "registrant_contact_name",
+                    "lobbyist_full_names",
+                    "previous_government_positions",
+                    "registrant_contact_telephone",
+                    "lobbyist_ids",
+                    "previous_goverment_positions"
+                
+            ]
+        },
+        "cfpb": {
+            "always": ["consumer_complaint_narrative"]
+        },
+        "clinical_trials": {
+            "always": ["responsible_party_investigator_name"]
+        }
+    },
+    "sovai": {
+        # Define sovai version exclusions here if needed
+        # For now, no exclusions for sovai version
+    }
+}
+
 # Check for full installation
 try:
    from sovai.extensions.pandas_extensions import CustomDataFrame
-   from sovai.utils.plot import plotting_data
+   from sovai.utils.plot import plotting_data 
+   import plotly.graph_objects as go
+
    HAS_FULL_INSTALL = True
 except ImportError:
    HAS_FULL_INSTALL = False
    CustomDataFrame = pd.DataFrame # Fallback if not available
    plotting_data = lambda x: None # Fallback
    logger.warning("Running with limited installation. Some features unavailable.")
+   
+   # Create a mock go module for type hints when full install is not available
+   if TYPE_CHECKING:
+       import plotly.graph_objects as go
 
 
 def is_full_installation() -> bool:
@@ -337,6 +372,7 @@ ENDPOINT_TO_TICKER = {
    "/breakout/median": "",
    "/institutional/trading": "",
    "/institutional/flow_prediction": "",
+   
    "/news/daily": "",
    "/news/match_quality": "",
    "/news/within_article": "",
@@ -353,6 +389,7 @@ ENDPOINT_TO_TICKER = {
    "/news/activeness": "",
    "/news/pronouns": "",
    "/news/word_count": "",
+
    "/news/sentiment_score": "", # was None before?
 
    "/insider/trading": "",
@@ -480,7 +517,24 @@ CLIENT_SIDE_ENDPOINTS_S3_PART_HIGH = {
     "short/maker",
     "short/over_shorted",
 
-    "institutional/trading"
+    "institutional/trading",
+
+    "news/daily",
+    "news/match_quality",
+    "news/within_article",
+    "news/relevance",
+    "news/magnitude",
+    "news/sentiment",
+    "news/article_count",
+    "news/associated_people",
+    "news/associated_companies",
+    "news/tone",
+    "news/positive",
+    "news/negative",
+    "news/polarity",
+    "news/activeness",
+    "news/pronouns",
+    "news/word_count"
 
 }
 
@@ -691,6 +745,93 @@ def read_parquet(url: str, use_polars: bool = False) -> DataFrameType:
 
 DEFAULT_COLUMNS_TO_REMOVE = ['ticker_partitioned', 'date_partitioned', 'year_partitioned', 'month_partitioned', 'day_partitioned']
 
+
+def detect_dataset_type(endpoint: str) -> Optional[str]:
+    """
+    Detect dataset type from endpoint URL for version-based column exclusions.
+    
+    Args:
+        endpoint: The API endpoint path
+        
+    Returns:
+        Dataset type string ('lobbying', 'cfpb', 'clinical_trials') or None
+    """
+    endpoint_lower = endpoint.lower().strip('/')
+    
+    if 'lobbying' in endpoint_lower:
+        return 'lobbying'
+    elif 'complaint' in endpoint_lower or 'cfpb' in endpoint_lower:
+        return 'cfpb'
+    elif 'clinical' in endpoint_lower or 'trial' in endpoint_lower:
+        return 'clinical_trials'
+    
+    return None
+
+
+def apply_version_exclusions(
+    data: DataFrameType,
+    version: Optional[str],
+    endpoint: str,
+    use_polars: bool = False
+) -> DataFrameType:
+    """
+    Apply version-based column exclusions to a DataFrame.
+    
+    Args:
+        data: Input DataFrame
+        version: Version parameter (e.g., 'p72', 'sovai')
+        endpoint: API endpoint to detect dataset type
+        use_polars: Whether using polars DataFrame
+        
+    Returns:
+        DataFrame with excluded columns removed
+    """
+    if data is None or (hasattr(data, 'empty') and data.empty):
+        return data
+    
+    if not version or version not in VERSION_COLUMN_EXCLUSIONS:
+        return data
+    
+    dataset_type = detect_dataset_type(endpoint)
+    if not dataset_type:
+        return data
+    
+    version_config = VERSION_COLUMN_EXCLUSIONS.get(version, {})
+    dataset_config = version_config.get(dataset_type, {})
+    
+    if not dataset_config:
+        return data
+    
+    # Get columns to exclude
+    columns_to_exclude = dataset_config.get('always', [])
+    
+    if not columns_to_exclude:
+        return data
+    
+    # Case-insensitive matching - find columns that exist in the DataFrame
+    data_columns = list(data.columns)
+    data_columns_lower = {col.lower(): col for col in data_columns}
+    
+    actual_columns_to_exclude = []
+    for exclude_col in columns_to_exclude:
+        exclude_col_lower = exclude_col.lower()
+        if exclude_col_lower in data_columns_lower:
+            actual_columns_to_exclude.append(data_columns_lower[exclude_col_lower])
+    
+    if not actual_columns_to_exclude:
+        verbose_mode.log(f"Version {version}: No matching columns found to exclude for {dataset_type}")
+        return data
+    
+    verbose_mode.log(f"Version {version}: Excluding columns {actual_columns_to_exclude} from {dataset_type}")
+    
+    # Remove the columns
+    if use_polars:
+        columns_to_keep = [col for col in data.columns if col not in actual_columns_to_exclude]
+        return data.select(columns_to_keep)
+    else:  # pandas
+        return data.drop(columns=actual_columns_to_exclude, errors='ignore')
+
+
 def filter_data(
    data: DataFrameType,
    columns: Optional[Union[str, List[str]]] = None,
@@ -1156,7 +1297,7 @@ def _uniform_datetime_params(datetime_params: Dict[str, str]) -> None:
                 logger.debug(f"Date parameter '{key}' with value '{original_val}' did not match any known formats: {datetime_format}. Leaving as is.")
 
 
-def _draw_graphs(data: Union[Dict, List[Dict]]) -> Optional[go.Figure]:
+def _draw_graphs(data: Union[Dict, List[Dict]]) -> Optional['go.Figure']:
     """
     Generate plots from data.
     
@@ -1206,7 +1347,7 @@ def _draw_graphs(data: Union[Dict, List[Dict]]) -> Optional[go.Figure]:
         return None
 
 
-def set_dark_mode(fig: go.Figure) -> go.Figure:
+def set_dark_mode(fig: 'go.Figure') -> 'go.Figure':
     """
     Apply dark mode styling to a Plotly figure.
     
@@ -1305,7 +1446,7 @@ def data(
     verbose: bool = False,
     full_history: bool = False,
     source: Optional[str] = None,
-) -> Union[DataFrameType, go.Figure, None]:
+) -> Union[DataFrameType, 'go.Figure', None]:
     """
     Main function to retrieve data from the API.
     Caches results to disk in 'cache/' directory (Parquet for DataFrames, Pickle for Figures).
@@ -1594,7 +1735,20 @@ def data(
             except Exception as err: verbose_mode.log(f"General API/processing error: {err}"); raise
     
     if result_data is not None:
-        result_data = process_dataframe(result_data) 
+        result_data = process_dataframe(result_data)
+        
+        # Apply version-based column exclusions if a version is set
+        # Use the version from function parameter, or fall back to ApiConfig.version
+        effective_version = version if version is not None else ApiConfig.version
+        
+        if effective_version and isinstance(result_data, (pd.DataFrame, pl.DataFrame)):
+            verbose_mode.log(f"Applying version-based exclusions for version: {effective_version}")
+            result_data = apply_version_exclusions(
+                result_data, 
+                effective_version, 
+                final_endpoint_path, 
+                use_polars
+            )
 
         if forced_columns_list is not None and isinstance(result_data, (pd.DataFrame, pl.DataFrame)):
             current_data_cols = list(result_data.columns)
@@ -1617,7 +1771,7 @@ async def fetch_multiple_tickers(
     endpoint: str,
     tickers: List[str],
     **kwargs
-) -> Dict[str, Union[DataFrameType, go.Figure, None]]: # Updated return type hint
+) -> Dict[str, Union[DataFrameType, 'go.Figure', None]]: # Updated return type hint
     """
     Fetch data for multiple tickers concurrently using asyncio.
     
@@ -1752,4 +1906,4 @@ ApiConfig.base_url = "your_api_base_url_here" # Ensure base_url is set
 
 # if __name__ == "__main__":
 #    asyncio.run(main())
-"""
+"""

{sovai-0.2.76 → sovai-0.2.77}/sovai/token_auth.py

@@ -1,16 +1,37 @@
+from typing import Optional
 from sovai.api_config import ApiConfig, save_key
 # Removed verify_token import from top level
 from sovai.errors.sovai_errors import InvalidCredentialsError
 
+# Valid version values
+VALID_VERSIONS = ["p72", "sovai"]
 
-def token_auth(token: str):
-    """Authenticates using a token, saves it, and verifies it."""
+
+def token_auth(token: str, version: Optional[str] = None):
+    """
+    Authenticates using a token, saves it, and verifies it.
+    
+    Args:
+        token: API authentication token
+        version: Optional version parameter for column exclusions (e.g., 'p72', 'sovai')
+    
+    Raises:
+        InvalidCredentialsError: If token is invalid or expired
+        ValueError: If version is provided but not in the list of valid versions
+    """
+    # Validate version if provided
+    if version is not None and version not in VALID_VERSIONS:
+        raise ValueError(
+            f"Invalid version '{version}'. Must be one of: {', '.join(VALID_VERSIONS)}"
+        )
 
     # Lazy load verify_token here, just before use
     from sovai.utils.client_side import verify_token
 
     ApiConfig.token = token
     ApiConfig.token_type = "Bearer"
+    ApiConfig.version = version
+    
     save_key() # Assumes save_key doesn't trigger heavy imports
 
     # Now call the imported function
@@ -22,4 +43,4 @@ def token_auth(token: str):
 
     # Optionally, you could return something useful, like the validity or user_id
     # print(f"Token verified successfully for user ID: {user_id}") # Or use logging
-    # return is_valid, user_id
+    # return is_valid, user_id

{sovai-0.2.76 → sovai-0.2.77}/sovai/utils/client_side_s3_part_high.py

@@ -875,10 +875,85 @@ ENDPOINT_CONFIG = {
         "date_path": "sovai/sovai-institutional/trading/partitioned/date",
         "has_year": False
     },
-
+    "news/match_quality": {
+        "ticker_path": "sovai/sovai-news/partitioned/match_quality/ticker",
+        "date_path": "sovai/sovai-news/partitioned/match_quality/date",
+        "has_year": False
+    },
+    "news/within_article": {
+        "ticker_path": "sovai/sovai-news/partitioned/within_article/ticker",
+        "date_path": "sovai/sovai-news/partitioned/within_article/date",
+        "has_year": False
+    },
+    "news/relevance": {
+        "ticker_path": "sovai/sovai-news/partitioned/relevance/ticker",
+        "date_path": "sovai/sovai-news/partitioned/relevance/date",
+        "has_year": False
+    },
+    "news/magnitude": {
+        "ticker_path": "sovai/sovai-news/partitioned/magnitude/ticker",
+        "date_path": "sovai/sovai-news/partitioned/magnitude/date",
+        "has_year": False
+    },
+    "news/sentiment": {
+        "ticker_path": "sovai/sovai-news/partitioned/sentiment/ticker",
+        "date_path": "sovai/sovai-news/partitioned/sentiment/date",
+        "has_year": False
+    },
+    "news/article_count": {
+        "ticker_path": "sovai/sovai-news/partitioned/article_count/ticker",
+        "date_path": "sovai/sovai-news/partitioned/article_count/date",
+        "has_year": False
+    },
+    "news/associated_people": {
+        "ticker_path": "sovai/sovai-news/partitioned/associated_people/ticker",
+        "date_path": "sovai/sovai-news/partitioned/associated_people/date",
+        "has_year": False
+    },
+    "news/associated_companies": {
+        "ticker_path": "sovai/sovai-news/partitioned/associated_companies/ticker",
+        "date_path": "sovai/sovai-news/partitioned/associated_companies/date",
+        "has_year": False
+    },
+    "news/tone": {
+        "ticker_path": "sovai/sovai-news/partitioned/tone/ticker",
+        "date_path": "sovai/sovai-news/partitioned/tone/date",
+        "has_year": False
+    },
+    "news/positive": {
+        "ticker_path": "sovai/sovai-news/partitioned/positive/ticker",
+        "date_path": "sovai/sovai-news/partitioned/positive/date",
+        "has_year": False
+    },
+    "news/negative": {
+        "ticker_path": "sovai/sovai-news/partitioned/negative/ticker",
+        "date_path": "sovai/sovai-news/partitioned/negative/date",
+        "has_year": False
+    },
+    "news/polarity": {
+        "ticker_path": "sovai/sovai-news/partitioned/polarity/ticker",
+        "date_path": "sovai/sovai-news/partitioned/polarity/date",
+        "has_year": False
+    },
+    "news/activeness": {
+        "ticker_path": "sovai/sovai-news/partitioned/activeness/ticker",
+        "date_path": "sovai/sovai-news/partitioned/activeness/date",
+        "has_year": False
+    },
+    "news/pronouns": {
+        "ticker_path": "sovai/sovai-news/partitioned/pronouns/ticker",
+        "date_path": "sovai/sovai-news/partitioned/pronouns/date",
+        "has_year": False
+    },
+    "news/word_count": {
+        "ticker_path": "sovai/sovai-news/partitioned/word_count/ticker",
+        "date_path": "sovai/sovai-news/partitioned/word_count/date",
+        "has_year": False
+    }
 }
 
 
+
 # =========================================================================
 # Public Interface
 # =========================================================================