bigdata-smart-batching 1.3.1__tar.gz → 1.4.0__tar.gz

{bigdata_smart_batching-1.3.1 → bigdata_smart_batching-1.4.0}/.gitignore

@@ -110,3 +110,5 @@ Thumbs.db
 *.csv
 
 benchmark/
+*.parquet
+*.json

{bigdata_smart_batching-1.3.1 → bigdata_smart_batching-1.4.0}/CHANGELOG.md

@@ -6,11 +6,17 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
-## [1.3.1] - 2026-05-18
+## [1.4.0] - 2026-05-20
 
 ### Added
 
-- Document-level **source** filters: `INCLUDE` / `EXCLUDE` with a non-empty list of source ID strings, validated with `validate_source_filter`, modeled as `SourceFilter` / `SourceFilterInput` with `SourceMode` (type alias of `CategoryMode`). Serialized to `query.filters.source` on co-mention, volume, and search payloads. `plan_search(..., source=...)` and `source_filter` plumbed through `SmartBatchingPlanner` co-mention and volume flows. Payload builders `build_comention_payload`, `build_volume_payload`, and `build_search_query` accept optional `source_filter`. Types and validator re-exported from the package root.
+- `execute_search(..., basket_filtered_entities=True)`: when enabled, each chunk's `entity_ids` list is filtered to only include entity IDs present in the basket's `query.filters.entity.any_of` list, removing co-mentioned entities that were not explicitly targeted by the basket. Entity IDs are always deduplicated.
+
+## [1.3.1] - 2026-05-19
+
+### Fixed
+
+- Now we always retry requests, before we were only retrying on 429 and 403 errors.
 
 ## [1.3.0] - 2026-05-12
 

{bigdata_smart_batching-1.3.1 → bigdata_smart_batching-1.4.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: bigdata-smart-batching
-Version: 1.3.1
+Version: 1.4.0
 Summary: High-performance semantic search with intelligent company grouping and parallel execution
 Project-URL: homepage, https://bigdata.com/api
 Author-email: "Bigdata.com" <support@ravenpack.com>
@@ -319,6 +319,7 @@ Step 2: EXECUTION
 | `max_workers` | `int` | `40` | Parallel workers |
 | `overwrite_chunks_per_basket` | `int \| None` | `None` | Override the proportional chunk allocation |
 | `second_pass` | `bool` | `False` | Re-query each basket for entities not yet retrieved |
+| `basket_filtered_entities` | `bool` | `False` | When `True`, each chunk's `entity_ids` is filtered to only include IDs present in the basket's `query.filters.entity.any_of` list; entity IDs are always deduplicated |
 
 ### Helper Functions
 
@@ -417,7 +418,7 @@ Ties are resolved using `VOLUME_BUCKET_PRIORITY`: `high > medium > low > very_lo
 ## Changelog
 
 See [CHANGELOG.md](CHANGELOG.md) for release notes. The package is currently at
-version **1.2.0**.
+version **1.4.0**.
 
 ## License
 

{bigdata_smart_batching-1.3.1 → bigdata_smart_batching-1.4.0}/README.md

@@ -303,6 +303,7 @@ Step 2: EXECUTION
 | `max_workers` | `int` | `40` | Parallel workers |
 | `overwrite_chunks_per_basket` | `int \| None` | `None` | Override the proportional chunk allocation |
 | `second_pass` | `bool` | `False` | Re-query each basket for entities not yet retrieved |
+| `basket_filtered_entities` | `bool` | `False` | When `True`, each chunk's `entity_ids` is filtered to only include IDs present in the basket's `query.filters.entity.any_of` list; entity IDs are always deduplicated |
 
 ### Helper Functions
 
@@ -401,7 +402,7 @@ Ties are resolved using `VOLUME_BUCKET_PRIORITY`: `high > medium > low > very_lo
 ## Changelog
 
 See [CHANGELOG.md](CHANGELOG.md) for release notes. The package is currently at
-version **1.2.0**.
+version **1.4.0**.
 
 ## License
 

{bigdata_smart_batching-1.3.1 → bigdata_smart_batching-1.4.0}/examples/example.py

@@ -45,6 +45,54 @@ def main() -> None:
         )
         sys.exit(1)
 
+    if False:
+        #Compare the results_raw parquet files to see if they are the same but just counting the number of ids
+        import pandas as pd
+        files = [x for x in os.listdir(".") if x.startswith("results_raw_") and x.endswith("_100pct_basket_6_without_sentiment.parquet")]
+        results_dfs = []
+        for file in files:
+            df = pd.read_parquet(file)
+            df_rows = []
+            for row in df.to_dict(orient="records"):
+                df_rows.append({
+                    'doc_id': row['id'],
+                    'chunk_cnum': [i['cnum'] for i in row['chunks']],
+                })
+            results_temp = pd.DataFrame(df_rows)
+            # Expand chunks into rows keeping document id, and chunk number
+            results_temp = results_temp.explode('chunk_cnum')
+            results_temp = results_temp.drop_duplicates()
+            #results_temp.sort_values(by=['doc_id', 'chunk_cnum']).to_csv(f"results_raw_{file.split('_')[2]}_100pct_basket_6.csv", index=False)
+            results_dfs.append(results_temp)                
+
+        # Merge using id as the key
+        merged = results_dfs[0].merge(results_dfs[1], on=['doc_id', 'chunk_cnum'], how='outer', indicator=True)
+        print(f"% overlapping chunks: {(len(merged[merged._merge!='both'])/2)/len(merged)*100:.2f}%")
+        print('a')        
+        
+
+    plan = load_plan("pricing_power_evidence.json")
+    # Set the query.ranking.params.content_diversification to False
+    for basket in plan['baskets']:
+        basket['query']['ranking_params']['content_diversification']['enabled'] = False
+        #Remove sentiment filter, do not set it to None, remove the key instead
+        basket['query']['filters'].pop('sentiment', None)
+    
+    # Keep only basket basket_6_high_20250924_20251022
+    plan['baskets'] = [basket for basket in plan['baskets'] if basket['basket_id'] == "basket_6_high_20250924_20251022"]
+    results_raw = execute_search(
+        search_plan=plan,
+        chunk_percentage=1,
+        requests_per_minute=450,
+        basket_filtered_entities=True,
+    )
+
+    # Store the results in a file but in parquet format
+    import json
+    import pandas as pd
+    df = pd.DataFrame(results_raw)
+    df.to_parquet(f"results_raw_{time.time()}_100pct_basket_6_without_sentiment.parquet")
+
     plans = []
     results_dfs = []
     for i in range(1):

bigdata_smart_batching-1.4.0/examples/example_non_determinism.py

@@ -0,0 +1,97 @@
+"""Execute a single Bigdata search request for basket_6_high_20250924_20251022.
+
+Uses only the requests library (no smart-batching helpers). Run twice and compare
+responses to observe search non-determinism.
+"""
+
+from __future__ import annotations
+
+import pandas as pd
+import requests
+
+SEARCH_ENDPOINT = "/v1/search"
+DEFAULT_API_BASE_URL = "https://api.bigdata.com"
+
+QUERY = {
+        "auto_enrich_filters": False,
+        "filters": {
+          "timestamp": {
+            "start": "2025-09-24T00:00:00Z",
+            "end": "2025-10-22T23:59:59Z"
+          },
+          "entity": {
+            "any_of": [
+              "12A3A3",
+              "619882",
+              "9196A2",
+              "58CA9A",
+              "50070E",
+              "E8846E",
+              "E09E2B",
+              "990AD0",
+              "4A6F00",
+              "0157B1"
+            ],
+            "search_in": "BODY"
+          },
+          "category": {
+            "mode": "INCLUDE",
+            "values": [
+              "news_premium",
+              "transcripts"
+            ]
+          },
+          "source": {
+            "mode": "EXCLUDE",
+            "values": [
+              "DFF004"
+            ]
+          },
+          "sentiment": {"ranges": [{"min": 0.1,"max": 1.0}]}
+        },
+        "ranking_params": {
+          "source_boost": 1,
+          "freshness_boost": 0,
+          "reranker": {
+            "enabled": False
+          },
+          "content_diversification": {"enabled": True}
+        },
+        "max_chunks": 1000,
+        "text": "The company demonstrates pricing power, raises prices without losing volume, or maintains price discipline against competitors."
+      }
+
+# Query from pricing_power_evidence.json plan, basket basket_6_high_20250924_20251022
+
+
+def main() -> None:
+
+    api_key = "bd_v1_UPdNGgq44wBO9XR8THfiGvHd7Km9TCIvDmnT5AiatA9T56"
+    url = "https://api.bigdata.com/v1/search"
+    headers = {"X-API-KEY": api_key, "Content-Type": "application/json"}
+    payload = {"query": QUERY}
+
+    # Execute the query 10 times and store the results in a df, counting the number of times a doc_id and chunk_cnum appears
+    dfs = []
+    for i in range(2):
+        df_rows = []
+        response = requests.post(url, headers=headers, json=payload, timeout=120)
+        results = response.json().get("results", [])
+        for result in results:
+            df_rows.append({
+                'doc_id': result['id'],
+                'chunk_cnum': [i['cnum'] for i in result['chunks']],
+            })
+        df = pd.DataFrame(df_rows)
+        # Expand chunks into rows keeping document id, and chunk number
+        df = df.explode('chunk_cnum')
+        df = df.drop_duplicates().reset_index(drop=True)
+        dfs.append(df)
+    
+    all_df = pd.concat(dfs)
+    all_df = all_df.groupby(['doc_id', 'chunk_cnum']).size().reset_index(name='count')
+    print(all_df)
+
+
+if __name__ == "__main__":
+    main()

{bigdata_smart_batching-1.3.1 → bigdata_smart_batching-1.4.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "bigdata-smart-batching"
-version = "1.3.1"
+version = "1.4.0"
 description = "High-performance semantic search with intelligent company grouping and parallel execution"
 readme = "README.md"
 authors = [{ name = "Bigdata.com", email = "support@ravenpack.com" }]

{bigdata_smart_batching-1.3.1 → bigdata_smart_batching-1.4.0}/src/bigdata_smart_batching/search_function.py

@@ -1247,6 +1247,7 @@ def execute_search(
     second_pass: bool = False,
     rate_limiter: SlidingWindowRateLimiter | None = None,
     concurrency_limiter: ConcurrencySemaphore | None = None,
+    basket_filtered_entities: bool = False,
 ) -> list[dict]:
     """
     Execute search with proportional sampling.
@@ -1270,6 +1271,9 @@ def execute_search(
         concurrency_limiter: Optional pre-built :class:`ConcurrencySemaphore` to
             share with other components. When provided, ``max_workers`` controls
             only the executor's worker count for basket fan-out.
+        basket_filtered_entities: If True, each chunk's ``entity_ids`` will be
+            filtered to only include entity IDs present in the basket's
+            ``query.filters.entity.any_of`` list.
 
     Returns:
         List of document dictionaries (each containing a chunks array).
@@ -1422,7 +1426,15 @@ def execute_search(
                                 entity_ids = [reporting_entities] if reporting_entities else []
 
                         # Store entity information (always set, even if empty)
-                        chunk["entity_ids"] = entity_ids
+                        if basket_filtered_entities:
+                            basket_entity_ids = set(
+                                basket_info.get("query", {})
+                                .get("filters", {})
+                                .get("entity", {})
+                                .get("any_of", [])
+                            )
+                            entity_ids = [eid for eid in entity_ids if eid in basket_entity_ids]
+                        chunk["entity_ids"] = list(set(entity_ids))
                         chunk["primary_entity_id"] = entity_ids[0] if entity_ids else None
 
                 logger.info(