datamule 0.422__cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

datamule/helper.py

@@ -0,0 +1,123 @@
+import requests
+import os
+from tqdm import tqdm
+import zipfile
+from pkg_resources import resource_filename
+import csv
+import re
+
+# Unused in current implementation.
+def construct_primary_doc_url(cik, accession_number,primary_doc_url):
+    accession_number = accession_number.replace("-", "")
+    return f"https://www.sec.gov/Archives/edgar/data/{cik}/{accession_number}/{primary_doc_url}"
+
+# DONE
+def _download_from_dropbox(url, output_path):
+    headers = {'user-agent': 'Wget/1.16 (linux-gnu)'}
+    r = requests.get(url, stream=True, headers=headers)
+    total_size = int(r.headers.get('content-length', 0))
+    
+    with open(output_path, 'wb') as f, tqdm(
+        desc="Downloading " + os.path.basename(output_path),
+        total=total_size,
+        unit='iB',
+        unit_scale=True,
+        unit_divisor=1024,
+    ) as progress_bar:
+        for chunk in r.iter_content(chunk_size=1024):
+            size = f.write(chunk)
+            progress_bar.update(size)
+    
+    # Check if the downloaded file is a zip file
+    if zipfile.is_zipfile(output_path):
+        extract_path = os.path.dirname(output_path)
+        with zipfile.ZipFile(output_path, 'r') as zip_ref:
+            for file_info in zip_ref.infolist():
+                extract_file_path = os.path.join(extract_path, file_info.filename)
+                with zip_ref.open(file_info) as file_in_zip, \
+                    open(extract_file_path, 'wb') as output_file, \
+                    tqdm(total=file_info.file_size, unit='B', unit_scale=True, 
+                         desc=f"Extracting {file_info.filename}") as pbar:
+                    while True:
+                        chunk = file_in_zip.read(8192)
+                        if not chunk:
+                            break
+                        output_file.write(chunk)
+                        pbar.update(len(chunk))
+        
+        # Remove the zip file after extraction
+        os.remove(output_path)
+        print(f"Extracted contents to {extract_path}")
+    else:
+        print(f"Downloaded file is not a zip. Saved to {output_path}")
+
+# May generalize to load any package resource
+def load_package_csv(name):
+    """Load package CSV files"""
+    csv_path = resource_filename('datamule', f'data/{name}.csv')
+    company_tickers = []
+    
+    with open(csv_path, 'r') as csvfile:
+        csv_reader = csv.DictReader(csvfile)
+        for row in csv_reader:
+            company_tickers.append(row)
+    
+    return company_tickers
+
+def load_package_dataset(dataset):
+    if dataset == 'company_tickers':
+        return load_package_csv('company_tickers')
+    elif dataset =='company_former_names':
+        return load_package_csv('company_former_names')
+    elif dataset =='company_metadata':
+        return load_package_csv('company_metadata')
+    elif dataset == 'sec_glossary':
+        return load_package_csv('sec-glossary')
+    elif dataset == 'xbrl_descriptions':
+        return load_package_csv('xbrl_descriptions')
+
+# DONE
+def identifier_to_cik(ticker):
+    """Convert company tickers to CIK codes"""
+    company_tickers = load_package_csv('company_tickers')
+    if ticker:
+        if isinstance(ticker, list):
+            cik = []
+            for t in ticker:
+                cik.extend([company['cik'] for company in company_tickers if t == company['ticker']])
+        else:
+            cik = [company['cik'] for company in company_tickers if ticker == company['ticker']]
+
+    if not cik:
+        raise ValueError("No matching companies found")
+
+    return cik
+
+
+def fix_filing_url(url):
+    match_suffix = re.search(r'/(\d{4})\.(.+?)$', url)
+    if match_suffix:
+        suffix_number = match_suffix.group(1)
+        file_ext = match_suffix.group(2)
+        match_accession = re.search(r'/(\d{18})/', url)
+        if match_accession:
+            accession_number = match_accession.group(1)
+            formatted_accession_number = f"{accession_number[:10]}-{accession_number[10:12]}-{accession_number[12:]}"
+            new_url = url.rsplit('/', 1)[0] + f'/{formatted_accession_number}-{suffix_number}.{file_ext}'
+            return new_url
+    return url
+
+def convert_to_dashed_accession(accession):
+    # Remove any existing dashes or whitespace
+    cleaned = ''.join(accession.split())
+    
+    # Check if the cleaned string has 18 characters
+    if len(cleaned) != 18:
+        raise ValueError("Invalid accession number format. Expected 18 characters.")
+    
+    # Insert dashes at the correct positions
+    dashed = f"{cleaned[:10]}-{cleaned[10:12]}-{cleaned[12:]}"
+    
+    return dashed
+
+headers = {'User-Agent': 'John Smith johnsmith@gmail.com'}

datamule/monitor.py

@@ -0,0 +1,236 @@
+import asyncio
+import aiohttp
+from datetime import timedelta, datetime
+import pytz
+from collections import deque
+import time
+from .helper import headers, identifier_to_cik
+
+def _get_current_eastern_date():
+    """Get current date in US Eastern timezone (automatically handles DST) """
+    eastern = pytz.timezone('America/New_York')
+    return datetime.now(eastern)
+
+class PreciseRateLimiter:
+    def __init__(self, rate, interval=1.0):
+        self.rate = rate  # requests per interval
+        self.interval = interval  # in seconds
+        self.token_time = self.interval / self.rate  # time per token
+        self.last_time = time.time()
+        self.lock = asyncio.Lock()
+
+    async def acquire(self):
+        async with self.lock:
+            now = time.time()
+            wait_time = self.last_time + self.token_time - now
+            if wait_time > 0:
+                await asyncio.sleep(wait_time)
+            self.last_time = time.time()
+            return True
+
+    async def __aenter__(self):
+        await self.acquire()
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        pass
+
+class RateMonitor:
+    def __init__(self, window_size=1.0):
+        self.window_size = window_size
+        self.requests = deque()
+        self._lock = asyncio.Lock()
+    
+    async def add_request(self, size_bytes):
+        async with self._lock:
+            now = time.time()
+            self.requests.append((now, size_bytes))
+            while self.requests and self.requests[0][0] < now - self.window_size:
+                self.requests.popleft()
+    
+    def get_current_rates(self):
+        now = time.time()
+        while self.requests and self.requests[0][0] < now - self.window_size:
+            self.requests.popleft()
+        
+        if not self.requests:
+            return 0, 0
+        
+        request_count = len(self.requests)
+        byte_count = sum(size for _, size in self.requests)
+        
+        requests_per_second = request_count / self.window_size
+        mb_per_second = (byte_count / 1024 / 1024) / self.window_size
+        
+        return round(requests_per_second, 1), round(mb_per_second, 2)
+
+class Monitor:
+    def __init__(self):
+        self.last_total = 0
+        self.last_date = _get_current_eastern_date()
+        self.submissions = []
+        self.max_hits = 10000
+        self.limiter = PreciseRateLimiter(5)  # 5 requests per second
+        self.rate_monitor = RateMonitor()
+        self.headers = headers
+
+    async def _fetch_json(self, session, url):
+        """Fetch JSON with rate limiting and monitoring."""
+        async with self.limiter:
+            try:
+                async with session.get(url) as response:
+                    response.raise_for_status()
+                    content = await response.read()
+                    await self.rate_monitor.add_request(len(content))
+                    return await response.json()
+            except Exception as e:
+                print(f"Error fetching {url}: {str(e)}")
+                return None
+
+    async def _poll(self, base_url, session, poll_interval, quiet):
+        """Poll API until new submissions are found."""
+        while True:
+            current_date = _get_current_eastern_date()
+            date_str = current_date.strftime('%Y-%m-%d')
+            
+            if self.last_date != current_date.strftime('%Y-%m-%d'):
+                print(f"New date: {date_str}")
+                self.last_total = 0
+                self.submissions = []
+                self.last_date = date_str
+            
+            poll_url = f"{base_url}&startdt={date_str}&enddt={date_str}"
+            if not quiet:
+                print(f"Polling {poll_url}")
+            
+            try:
+                data = await self._fetch_json(session, poll_url)
+                if data:
+                    current_total = data['hits']['total']['value']
+                    if current_total > self.last_total:
+                        print(f"Found {current_total - self.last_total} new submissions")
+                        self.last_total = current_total
+                        return current_total, data, poll_url
+                    self.last_total = current_total
+            except Exception as e:
+                print(f"Error in poll: {str(e)}")
+            
+            await asyncio.sleep(poll_interval / 1000)
+
+    async def _retrieve_batch(self, session, poll_url, from_positions, quiet):
+        """Retrieve a batch of submissions concurrently."""
+        tasks = [
+            self._fetch_json(
+                session,
+                f"{poll_url}&from={pos}"
+            )
+            for pos in from_positions
+        ]
+        
+        results = await asyncio.gather(*tasks, return_exceptions=True)
+        submissions = []
+        
+        for result in results:
+            if isinstance(result, Exception):
+                print(f"Error in batch: {str(result)}")
+                continue
+            if result and 'hits' in result:
+                submissions.extend(result['hits']['hits'])
+        
+        return submissions
+
+    async def _retrieve(self, poll_url, initial_data, session, quiet):
+        """Retrieve all submissions using parallel batch processing."""
+        batch_size = 10  # Number of concurrent requests
+        page_size = 100  # Results per request
+        max_position = min(self.max_hits, self.last_total)
+        submissions = []
+        
+        # Process in batches of concurrent requests
+        for batch_start in range(0, max_position, batch_size * page_size):
+            from_positions = [
+                pos for pos in range(
+                    batch_start,
+                    min(batch_start + batch_size * page_size, max_position),
+                    page_size
+                )
+            ]
+            
+            if not quiet:
+                print(f"Retrieving batch from positions: {from_positions}")
+            
+            batch_submissions = await self._retrieve_batch(
+                session, poll_url, from_positions, quiet
+            )
+            
+            if not batch_submissions:
+                break
+                
+            submissions.extend(batch_submissions)
+            
+            # If we got fewer results than expected, we're done
+            if len(batch_submissions) < len(from_positions) * page_size:
+                break
+        
+        return submissions
+
+    async def _monitor(self, callback, form=None, cik=None, ticker=None, poll_interval=1000, quiet=True):
+        """Main monitoring loop with parallel processing."""
+        if poll_interval < 100:
+            raise ValueError("SEC rate limit is 10 requests per second, set poll_interval to 100ms or higher")
+
+        # Handle form parameter
+        if form is None:
+            form = ['-0']
+        elif isinstance(form, str):
+            form = [form]
+        
+        # Handle CIK/ticker parameter
+        cik_param = None
+        if ticker is not None:
+            cik_param = identifier_to_cik(ticker)
+        elif cik is not None:
+            cik_param = cik if isinstance(cik, list) else [cik]
+
+        # Construct base URL
+        base_url = 'https://efts.sec.gov/LATEST/search-index?forms=' + ','.join(form)
+        
+        # Add CIK parameter if specified
+        if cik_param:
+            cik_list = ','.join(str(c).zfill(10) for c in cik_param)
+            base_url += f"&ciks={cik_list}"
+        
+        async with aiohttp.ClientSession(headers=self.headers) as session:
+            while True:
+                try:
+                    # Poll until we find new submissions
+                    _, data, poll_url = await self._poll(base_url, session, poll_interval, quiet)
+                    
+                    # Retrieve all submissions in parallel
+                    submissions = await self._retrieve(poll_url, data, session, quiet)
+                    
+                    # Find new submissions
+                    existing_ids = {sub['_id'] for sub in self.submissions}
+                    new_submissions = [
+                        sub for sub in submissions 
+                        if sub['_id'] not in existing_ids
+                    ]
+                    
+                    if new_submissions:
+                        self.submissions.extend(new_submissions)
+                        if callback:
+                            await callback(new_submissions)
+                        
+                        reqs_per_sec, mb_per_sec = self.rate_monitor.get_current_rates()
+                        if not quiet:
+                            print(f"Current rates: {reqs_per_sec} req/s, {mb_per_sec} MB/s")
+                    
+                except Exception as e:
+                    print(f"Error in monitor: {str(e)}")
+                    await asyncio.sleep(poll_interval / 1000)
+                
+                await asyncio.sleep(poll_interval / 1000)
+
+    def monitor_submissions(self, callback=None, form=None, cik=None, ticker=None, poll_interval=1000, quiet=True):
+        """Start the monitoring process."""
+        asyncio.run(self._monitor(callback, form, cik, ticker, poll_interval, quiet))

datamule/mulebot/__init__.py

@@ -0,0 +1 @@
+from .mulebot import MuleBot

datamule/mulebot/helper.py

@@ -0,0 +1,35 @@
+import requests
+from datamule.global_vars import headers
+from datamule.helper import identifier_to_cik
+from datamule import Parser
+
+parser = Parser()
+
+def get_company_concept(ticker):
+
+    cik = identifier_to_cik(ticker)[0]
+    url = f'https://data.sec.gov/api/xbrl/companyfacts/CIK{str(cik).zfill(10)}.json'
+    response = requests.get(url,headers=headers)
+    data = response.json()
+
+    table_dict_list = parser.parse_company_concepts(data)
+
+    # drop tables where label is None
+    table_dict_list = [table_dict for table_dict in table_dict_list if table_dict['label'] is not None]
+    
+    return table_dict_list
+
+def select_dict_by_title(data, title):
+    if isinstance(data, dict):
+        if data.get('title') == title:
+            return data
+        for value in data.values():
+            result = select_dict_by_title(value, title)
+            if result:
+                return result
+    elif isinstance(data, list):
+        for item in data:
+            result = select_dict_by_title(item, title)
+            if result:
+                return result
+    return None

datamule/mulebot/mulebot.py

@@ -0,0 +1,130 @@
+import openai
+import json
+
+from datamule.helper import identifier_to_cik
+from datamule import Downloader, Parser
+from .search import search_filing
+from .tools import tools, return_title_tool
+from .helper import get_company_concept, select_dict_by_title
+
+downloader = Downloader()
+parser = Parser()
+
+
+class MuleBot:
+    def __init__(self, api_key):
+        self.client = openai.OpenAI(api_key=api_key)
+        self.messages = [
+            {"role": "system", "content": "You are a helpful, but concise, assistant to assist with questions related to the Securities and Exchanges Commission. You are allowed to guess tickers."}
+        ]
+        self.total_tokens = 0
+
+    def process_message(self, user_input):
+
+        new_message_chain = self.messages
+        new_message_chain.append({"role": "user", "content": user_input})
+
+        try:
+            response = self.client.chat.completions.create(
+                model="gpt-4o-mini",
+                messages=new_message_chain,
+                tools=tools,
+                tool_choice="auto"
+            )
+
+            self.total_tokens += response.usage.total_tokens
+            assistant_message = response.choices[0].message
+
+            if assistant_message.content is None:
+                assistant_message.content = "I'm processing your request."
+
+            new_message_chain.append({"role": "assistant", "content": assistant_message.content})
+            
+            tool_calls = assistant_message.tool_calls
+            if tool_calls is None:
+                return {'key':'text','value':assistant_message.content}
+            else:
+                for tool_call in tool_calls:
+                    print(f"Tool call: {tool_call.function.name}")
+                    if tool_call.function.name == "identifier_to_cik":
+                        function_args = json.loads(tool_call.function.arguments)
+                        print(f"Function args: {function_args}")
+                        
+                        cik = identifier_to_cik(function_args["ticker"])
+                        return {'key':'text','value':cik}
+                    elif tool_call.function.name == "get_company_concept":
+                        function_args = json.loads(tool_call.function.arguments)
+                        print(f"Function args: {function_args}")
+                        table_dict_list = get_company_concept(function_args["ticker"])
+                        return {'key':'table','value':table_dict_list}
+                    elif tool_call.function.name == "get_filing_urls":
+                        function_args = json.loads(tool_call.function.arguments)
+                        print(f"Function args: {function_args}")
+                        result = downloader.download(**function_args,return_urls=True)
+                        return {'key':'list','value':result}
+                    elif tool_call.function.name == "find_filing_section_by_title":
+                        function_args = json.loads(tool_call.function.arguments)
+                        print(f"Function args: {function_args}")
+                        # Parse the filing
+                        data = parser.parse_filing(function_args["url"])
+
+                        # find possible matches
+                        section_dicts = search_filing(query = function_args["title"], nested_dict =data, score_cutoff=0.3)
+
+                        # feed titles back to assistant
+                        titles = [section['title'] for section in section_dicts]
+                        new_message_chain.append({"role": "assistant", "content": f"Which of these titles is closest: {','.join(titles)}"})
+
+                        title_response = self.client.chat.completions.create(
+                            model="gpt-4o-mini",
+                            messages=new_message_chain,
+                            tools=[return_title_tool],
+                            tool_choice="required"
+                        )
+
+                        title_tool_call = title_response.choices[0].message.tool_calls[0]
+                        title = json.loads(title_tool_call.function.arguments)['title']
+                        print(f"Selected title: {title}")
+                        #print(f"Possible titles: {titles}")
+
+                        # select the section
+                        #section_dict = select_dict_by_title(data, title)
+                        
+                        # probably want to return full dict, and section label
+                        return {'key':'filing','value':{'data':data,'title':title}}
+
+            return {'key':'text','value':'No tool call was made.'}
+
+        except Exception as e:
+            return f"An error occurred: {str(e)}"
+
+    def get_total_tokens(self):
+        return self.total_tokens
+    
+    def run(self):
+        """Basic chatbot loop"""
+        print("MuleBot: Hello! I'm here to assist you with questions related to the Securities and Exchange Commission. Type 'quit', 'exit', or 'bye' to end the conversation.")
+        while True:
+            user_input = input("You: ")
+            if user_input.lower() in ['quit', 'exit', 'bye']:
+                print("MuleBot: Goodbye!")
+                break
+            
+            response = self.process_message(user_input)
+            response_type = response['key']
+
+            if response_type == 'text':
+                value = response['value']
+                print(value)
+            elif response_type == 'table':
+                value = response['value']
+                print(value)
+            elif response_type == 'list':
+                value = response['value']
+                print(value)
+            elif response_type == 'filing':
+                value = response['value']
+                print(value)
+            else:
+                value = response['value']
+                print(value)

datamule/mulebot/mulebot_server/__init__.py

@@ -0,0 +1 @@
+from .server import MuleBotServer

datamule/mulebot/mulebot_server/server.py

@@ -0,0 +1,87 @@
+import os
+from flask import Flask, request, jsonify, render_template
+from datamule.mulebot import MuleBot
+from datamule.filing_viewer import create_interactive_filing, create_valid_id
+
+class MuleBotServer:
+    def __init__(self, template='chat-minimalist.html'):
+        template_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), 'templates'))
+        static_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), 'static'))
+        self.app = Flask(__name__, template_folder=template_dir, static_folder=static_dir)
+        self.mulebot = None
+        self.template = template
+        self.setup_routes()
+
+    def setup_routes(self):
+        @self.app.route('/')
+        def home():
+            return render_template(self.template)
+
+        @self.app.route('/chat-with-prompt')
+        def chat_with_prompt():
+            prefilled_prompt = request.args.get('prompt', '')
+            return render_template(self.template, prefilled_prompt=prefilled_prompt)
+
+        @self.app.route('/chat', methods=['POST'])
+        def chat():
+            user_input = request.json['message']
+            
+            # Process the message using MuleBot's process_message method
+            response = self.mulebot.process_message(user_input)
+            response_type = response['key']
+
+            # Prepare the response based on the type
+            if response_type == 'text':
+                # If response type is text, add it to the chat
+                chat_response = {
+                    'type': 'text',
+                    'content': response['value']
+                }
+            elif response_type == 'table':
+                # If response type is table, prepare it for the artifact window
+                chat_response = {
+                    'type': 'artifact',
+                    'content': response['value'],
+                    'artifact_type': 'artifact-table'
+                }
+            elif response_type == 'list':
+                chat_response = {
+                    'type': 'artifact',
+                    'content': response['value'],
+                    'artifact_type': 'artifact-list'
+                }
+            elif response_type == 'filing':
+                data = response['value']['data']
+                title = response['value']['title']
+                section_id = create_valid_id(title)
+
+                # create a filing viewer display
+                html = create_interactive_filing(data)
+                
+                # we'll need to display the filing viewer in the artifact window, with a json export option
+                chat_response = {
+                    'type': 'artifact',
+                    'content': html,
+                    'data': data,
+                    'section_id': section_id,
+                    'artifact_type': 'artifact-filing'
+                }
+            else:
+                # Handle other types of responses if needed
+                chat_response = {
+                    'type': 'unknown',
+                    'content': 'Unsupported response type'
+                }
+            
+            return jsonify({
+                'response': chat_response,
+                'total_tokens': self.mulebot.get_total_tokens()
+            })
+
+    def set_api_key(self, api_key):
+        self.mulebot = MuleBot(api_key)
+
+    def run(self, debug=False, host='0.0.0.0', port=5000):
+        if not self.mulebot:
+            raise ValueError("API key not set. Please call set_api_key() before running the server.")
+        self.app.run(debug=debug, host=host, port=port)