orbitkit 0.8.36__tar.gz → 0.8.38__tar.gz

{orbitkit-0.8.36/orbitkit.egg-info → orbitkit-0.8.38}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.1
 Name: orbitkit
-Version: 0.8.36
+Version: 0.8.38
 Summary: This project is only for Orbit Tech internal use.
 Home-page: https://github.com/clown-0726/orbitkit
 Author: Lilu Cao
@@ -28,18 +28,6 @@ Requires-Dist: prettytable>=3.0.0
 Requires-Dist: pytz>=2022.1
 Requires-Dist: Deprecated
 Requires-Dist: func_timeout
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: license
-Dynamic: maintainer
-Dynamic: maintainer-email
-Dynamic: platform
-Dynamic: requires-dist
-Dynamic: summary
 
 # orbitkit
 

orbitkit-0.8.38/orbitkit/VERSION

@@ -0,0 +1 @@
+0.8.38

{orbitkit-0.8.36 → orbitkit-0.8.38}/orbitkit/airflow_handler/data_preprocessing.py

@@ -1,3 +1,4 @@
+import os
 import datetime
 from collections import defaultdict
 
@@ -101,20 +102,35 @@ class DocumentProcessor:
         report_type_ids = doc.get('x_orbit_data', {}).get('report_type_id_list', [])
         return "extract" if report_type_ids == ['19999'] else target_stage
 
-    @staticmethod
-    def create_record(doc, start_stage):
+    @classmethod
+    def create_record(cls, doc, start_stage):
         attachments = doc.get('x_attachments', []) if start_stage == 'convert' else doc.get('x_attachments_pdf', [])
         s3_path_info = []
+        add_extends = {}
         for att in attachments:
             if len(att['file_name']) > 2000 or len(att['file_name'].encode('utf-8')) > 2000:
                 return False, None
+
+            if start_stage == 'convert' and not add_extends:
+                _, ext = os.path.splitext(att['store_path'])
+                if ext in cls.AUDIO_SUFFIXES or ext in cls.VIDEO_SUFFIXES:
+                    add_extends = {
+                        "title": doc['x_orbit_data']['report_title'],
+                        "published": doc['x_reported_at_utc_date'],
+                        "tickers": [],
+                        "perm_id_list": doc['x_orbit_data']['perm_id_list'],
+                        "report_type_id_list_str": doc['x_orbit_data']['report_type_id_list']
+                    }
+
             s3_path_info.append({
                 'store_path': f"s3://{att['bucket']}/{att['store_path']}" if start_stage == 'convert' else att[
                     'store_path'],
                 'file_name': att['file_name']
             })
-
-        return True, {'id': doc['_id'], 's3_path_info': s3_path_info}
+        result_dict = {'id': doc['_id'], 's3_path_info': s3_path_info}
+        if add_extends:
+            result_dict['extends'] = add_extends
+        return True, result_dict
 
     @staticmethod
     def create_result_info(process_type, message, result_data):

{orbitkit-0.8.36 → orbitkit-0.8.38}/orbitkit/airflow_handler/file_flow_entry_process.py

@@ -5,31 +5,40 @@ from typing import Literal
 import logging
 import pymongo
 import pytz
+import boto3
+from sqlalchemy import create_engine, Table, MetaData, select
+from sqlalchemy.orm import sessionmaker, scoped_session
+from contextlib import contextmanager
+
 from orbitkit.airflow_handler.file_handler_v2 import FileFlowHandleV2
 from orbitkit.airflow_handler.data_preprocessing import DocumentProcessor
+from orbitkit.orbit_type import OrbitTypeMatcher
 
 logger = logging.getLogger(__name__)
 
 
 class FilingOfficialProcessor:
 
-    def __init__(self, mongo_uri=None, postgres_uri=None):
+    def __init__(self, mongo_uri=None, postgres_uri=None, aws_access_key_id=None, aws_secret_access_key=None, pi2_postgres_uri=None, pi2_database_name=None):
         mongo_uri = os.environ.get('MONGO_URI_MAIN_USER_APP') if not mongo_uri else mongo_uri
         if not mongo_uri:
             raise KeyError('mongo_uri not set.')
 
+        if not aws_secret_access_key or not aws_access_key_id:
+            raise KeyError('aws_access_key_id and aws_secret_access_key not set.')
+
         self.mongo_client = pymongo.MongoClient(mongo_uri)
         self.data_xbrl_convert_collection = self.mongo_client['filing_reports']['data_xbrl_convert']
         self.filing_data_collection = self.mongo_client['filing_reports']['filing_data']
-        self.official_data_collection = self.mongo_client['filing_reports']['official_data_relocation']
+        self.filing_reports_astock_test0822_collection = self.mongo_client['filing_reports']['filing_reports_astock_test0822']
         self.annotation_reports_view_rows_collection = self.mongo_client['filing_reports'][
             'annotation_reports_view_rows']
         self.source_map = {
             'filing_data': (self.filing_data_collection, 'filing_data'),
-            'official_data': (self.official_data_collection, 'official_data_relocation'),
+            'filing_reports_astock_test0822': (self.filing_reports_astock_test0822_collection, 'filing_reports_astock_test0822'),
             'reports_view': [
                 (self.filing_data_collection, 'filing_data'),
-                (self.official_data_collection, 'official_data_relocation')
+                (self.filing_reports_astock_test0822_collection, 'filing_reports_astock_test0822')
             ]
         }
         self.only_low_important_set = {'internal_seekingalpha'}
@@ -41,6 +50,36 @@ class FilingOfficialProcessor:
         self.max_batch_size = 10000
         self.all_stat_count = {'all': 0, 'skip': 0, 'step_error': 0, 'xbrl': 0, 'file_flow': 0}
 
+        self.s3_client = boto3.client('s3', aws_access_key_id=aws_access_key_id, aws_secret_access_key=aws_secret_access_key)
+        self.matcher = OrbitTypeMatcher(self.s3_client)
+        self.report_type_id_name_map = {i["lv3_id"]: i["lv3_name"] for i in self.matcher.get_full_type_list()}
+
+        self.pi2_postgres_uri = pi2_postgres_uri or os.environ['PG_URI_CX45_USER_GLAUUIADMIN']
+        if not self.pi2_postgres_uri:
+            raise KeyError('pie_postgres_uri not set.')
+        self.databases = pi2_database_name or 'newsfeedsite'
+        self.postgres_engine = create_engine(f"{self.pi2_postgres_uri}/{self.databases}", connect_args={"sslmode": "require"})
+        self.postgres_session = sessionmaker(bind=self.postgres_engine)
+        self.Session = scoped_session(self.postgres_session)
+        self.postgres_metadata = MetaData()
+
+        self.pi2_table = Table(
+            'primary_instrument_2_release', self.postgres_metadata,
+            autoload_with=self.postgres_engine, schema='security_master'
+        )
+
+    @contextmanager
+    def session_scope(self):
+        session = self.Session()
+        try:
+            yield session
+            session.commit()
+        except Exception:
+            session.rollback()
+            raise
+        finally:
+            self.Session.remove()
+
     def create_spider_name_source_type_map(self, collections):
 
         def find_duplicates(keys):
@@ -85,6 +124,29 @@ class FilingOfficialProcessor:
         }})
         logger.info(f'The document file type cannot be converted.')
 
+    def update_extends_fields(self, perm_id_list, file_flow_info):
+        stmt = select(self.pi2_table.c.orbit_entity_id, self.pi2_table.c.ticker).where(self.pi2_table.c.orbit_entity_id.in_(perm_id_list))
+        orbit_entity_id_ticker_map = {}
+        with self.session_scope() as session:
+            result = session.execute(stmt)
+            for row in result:
+                if row.orbit_entity_id not in orbit_entity_id_ticker_map:
+                    orbit_entity_id_ticker_map[row.orbit_entity_id] = []
+
+                if row.ticker is not None:
+                    orbit_entity_id_ticker_map[row.orbit_entity_id].append(row.ticker)
+        for step_info, records in file_flow_info.items():
+            for record in records:
+                if 'extends' in record and record.get('extends') is not None:
+                    tickers = []
+                    for i in record['extends']['perm_id_list']:
+                        tickers.extend(orbit_entity_id_ticker_map.get(i, []))
+                    record['extends']['tickers'] = tickers
+
+                    record['extends']['report_type_id_list_str'] = [self.report_type_id_name_map.get(i) for i in record['extends']['report_type_id_list_str']]
+
+        return file_flow_info
+
     def send_task(self, file_flow_info, tags, is_important, priority, spider_name_source_type):
         for step_str, records in file_flow_info.items():
             steps = step_str.split('@__@')
@@ -116,7 +178,7 @@ class FilingOfficialProcessor:
             logger.info(f"{len(records)}--{start_stage}-{target_stage}-{x_spider_name} status: {status}, message: {message}")
 
 
-    def process_task_entry(self, source: Literal["filing_data", "official_data", "reports_view"],
+    def process_task_entry(self, source: Literal["filing_data", "filing_reports_astock_test0822", "reports_view"],
                            query: dict, tags: list[str], priority: str,
                            is_important: bool = False):
 
@@ -128,16 +190,20 @@ class FilingOfficialProcessor:
         spider_name_source_type = self.create_spider_name_source_type_map(collections)
 
         process_data = []
+        perm_id_set = set()
         for collection, label in collections:
             logger.info(f"load {label} data.")
             docs = collection.find(query).batch_size(1000)
 
             for doc in docs:
                 self.all_stat_count['all'] += 1
+                for orbit_entity_id in doc['x_orbit_data']['perm_id_list']:
+                    perm_id_set.add(orbit_entity_id)
                 process_data.append(self.data_processor.process(doc))
                 if len(process_data) >= self.max_batch_size:
                     file_flow_info, xbrl_data, except_id_list, doc_error_list = self.data_processor.split_data_by_spider_name_and_step(
                         process_data)
+                    file_flow_info = self.update_extends_fields(list(perm_id_set), file_flow_info)
                     self.all_stat_count['skip'] += len(doc_error_list)
                     self.all_stat_count['step_error'] += len(except_id_list)
                     self.all_stat_count['xbrl'] += len(xbrl_data)
@@ -145,10 +211,12 @@ class FilingOfficialProcessor:
                     self.send_xbrl_data_to_mongo(xbrl_data)
                     self.update_doc_status_to_convert(collection, doc_error_list)
                     process_data.clear()
+                    perm_id_set.clear()
 
             if process_data:
                 file_flow_info, xbrl_data, except_id_list, doc_error_list = self.data_processor.split_data_by_spider_name_and_step(
                     process_data)
+                file_flow_info = self.update_extends_fields(list(perm_id_set), file_flow_info)
                 self.all_stat_count['skip'] += len(doc_error_list)
                 self.all_stat_count['step_error'] += len(except_id_list)
                 self.all_stat_count['xbrl'] += len(xbrl_data)
@@ -156,5 +224,6 @@ class FilingOfficialProcessor:
                 self.send_xbrl_data_to_mongo(xbrl_data)
                 self.update_doc_status_to_convert(collection, doc_error_list)
                 process_data.clear()
+                perm_id_set.clear()
 
         logger.info(f"finish processing {self.all_stat_count}.")

{orbitkit-0.8.36 → orbitkit-0.8.38}/orbitkit/airflow_handler/file_handler_v2.py

@@ -132,6 +132,7 @@ class FileFlowHandleV2:
             'current_stage': params['current_stage'],
             'target_stage': params['target_stage'],
             'data_source': params['source_type'],
+            'extends': record.get('extends', {}),
             'created_at': now,
             'updated_at': now,
             'tags': params['tags'],

{orbitkit-0.8.36 → orbitkit-0.8.38}/orbitkit/util/__init__.py

@@ -67,3 +67,6 @@ from .util_url import (
     website_url_standardize,
     website_url_compare,
 )
+from .cache_asset_downloader import (
+    CacheAssetDownloader,
+)

orbitkit-0.8.38/orbitkit/util/cache_asset_downloader.py

@@ -0,0 +1,132 @@
+import os
+import platform
+import hashlib
+import requests
+from pathlib import Path
+from typing import Optional
+from urllib.parse import urlparse
+
+"""
+功能描述
+    CacheAssetDownloader 是一个智能的资源下载器，主要用于管理大型文件（如机器学习模型、数据集等）的下载和本地缓存。
+
+跨平台缓存管理：
+    Windows：使用 %LOCALAPPDATA% 或 用户目录/AppData/Local
+    macOS：使用 ~/Library/Caches
+    Linux/Docker：使用 $XDG_CACHE_HOME 或 ~/.cache
+
+智能缓存机制：
+    自动检查本地缓存是否存在
+    支持 MD5 哈希验证确保文件完整性
+    避免重复下载，节省带宽和时间
+
+安全下载：
+    流式下载，适合大文件
+    异常处理和错误恢复
+    文件验证失败时自动清理
+"""
+
+
+class CacheAssetDownloader:
+    def __init__(self, asset_sub_folder: str, download_url: str, expected_hash: Optional[str] = None):
+        self.asset_sub_folder = asset_sub_folder
+        self.download_url = download_url
+        self.expected_hash = expected_hash
+
+        # 设置缓存目录
+        self.cache_dir = self._get_cache_dir()
+        self.asset_path = None
+
+        # 自动下载资源
+        self._download_if_needed()
+
+    def _get_cache_dir(self) -> Path:
+        """获取缓存目录"""
+        system = platform.system().lower()
+
+        if system == 'windows':
+            base_dir = Path(os.environ.get('LOCALAPPDATA', Path.home() / 'AppData' / 'Local'))
+        elif system == 'darwin':  # macOS
+            base_dir = Path.home() / 'Library' / 'Caches'
+        else:  # Linux/Docker
+            base_dir = Path(os.environ.get('XDG_CACHE_HOME', Path.home() / '.cache'))
+
+        cache_dir = base_dir / 'orbit_assets' / self.asset_sub_folder
+        cache_dir.mkdir(parents=True, exist_ok=True)
+
+        print(f"You current system: >>>{system}<<<, with cache path: >>>{cache_dir}<<<")
+        return cache_dir
+
+    def _validate_file(self, file_path: Path) -> bool:
+        """验证文件"""
+        if not file_path.exists():
+            return False
+
+        if not self.expected_hash:  # 如果没开启 hash 验证，则直接返回
+            return True
+
+        # 计算MD5
+        hash_md5 = hashlib.md5()
+        with open(file_path, 'rb') as f:
+            for chunk in iter(lambda: f.read(4096), b""):
+                hash_md5.update(chunk)
+
+        return hash_md5.hexdigest().lower() == self.expected_hash.lower()
+
+    def _download_file(self, url: str, target_path: Path) -> bool:
+        """下载文件"""
+        try:
+            print(f"正在下载: {url}")
+            response = requests.get(url, stream=True)
+            response.raise_for_status()
+
+            with open(target_path, 'wb') as f:
+                for chunk in response.iter_content(chunk_size=8192):
+                    if chunk:
+                        f.write(chunk)
+
+            print(f"下载完成: {target_path}")
+            return True
+        except Exception as e:
+            print(f"下载失败: {e}")
+            return False
+
+    def _download_if_needed(self):
+        """检查并下载资源"""
+        # 确定文件名
+        filename = os.path.basename(urlparse(self.download_url).path)
+        asset_file = self.cache_dir / filename
+
+        # 检查缓存
+        if self._validate_file(asset_file):
+            print(f"使用缓存文件: {asset_file}")
+            self.asset_path = asset_file
+            return
+
+        # 下载文件
+        if not self._download_file(self.download_url, asset_file):
+            raise RuntimeError(f"下载失败: {self.download_url}")
+
+        # 验证文件
+        if not self._validate_file(asset_file):
+            asset_file.unlink(missing_ok=True)
+            raise RuntimeError("文件验证失败")
+
+        self.asset_path = asset_file
+
+    def get_path(self) -> Path:
+        """获取资源路径"""
+        return self.asset_path
+
+
+if __name__ == "__main__":
+    downloader = CacheAssetDownloader(
+        asset_sub_folder="fasttext",
+        download_url="https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin",
+    )
+
+    # cache_dir = downloader._get_cache_dir()
+    # print(cache_dir)
+
+    model_path = downloader.get_path()
+    print(f"模型路径: {model_path}")

{orbitkit-0.8.36 → orbitkit-0.8.38/orbitkit.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.1
 Name: orbitkit
-Version: 0.8.36
+Version: 0.8.38
 Summary: This project is only for Orbit Tech internal use.
 Home-page: https://github.com/clown-0726/orbitkit
 Author: Lilu Cao
@@ -28,18 +28,6 @@ Requires-Dist: prettytable>=3.0.0
 Requires-Dist: pytz>=2022.1
 Requires-Dist: Deprecated
 Requires-Dist: func_timeout
-Dynamic: author
-Dynamic: author-email
-Dynamic: classifier
-Dynamic: description
-Dynamic: description-content-type
-Dynamic: home-page
-Dynamic: license
-Dynamic: maintainer
-Dynamic: maintainer-email
-Dynamic: platform
-Dynamic: requires-dist
-Dynamic: summary
 
 # orbitkit
 

{orbitkit-0.8.36 → orbitkit-0.8.38}/orbitkit.egg-info/SOURCES.txt

@@ -56,6 +56,7 @@ orbitkit/pdf_extractor_simple/utils.py
 orbitkit/pdf_writer/__init__.py
 orbitkit/pdf_writer/pdf_writer_simple.py
 orbitkit/util/__init__.py
+orbitkit/util/cache_asset_downloader.py
 orbitkit/util/common.py
 orbitkit/util/customize_regix_manager.py
 orbitkit/util/secret_manager.py

orbitkit-0.8.36/orbitkit/VERSION

@@ -1 +0,0 @@
-0.8.36