Análisis de datos de acciones

import base64
import json
import os
from typing import Dict, List, Optional
import uuid
import urllib.request
import yfinance as yf
from datetime import datetime, timedelta

HOST = "your_zylon_instance.com"
TOKEN = "your api token here"

API_HOST = f"https://{HOST}/api/gpt"
VECTOR_COLLECTION = "stock_documents"


def http_post(url: str, data: dict, token=TOKEN) -> dict:
    """
    Simple POST request.

    All requests to Zylon API require an Authorization header with a Bearer token.
    """
    req = urllib.request.Request(
        url,
        data=json.dumps(data).encode("utf-8"),
        headers={
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json"
        },
        method="POST"
    )
    with urllib.request.urlopen(req, timeout=300) as resp:
        if resp.status != 200:
            raise Exception(f"HTTP {resp.status}: {resp.read().decode('utf-8')}")
        return json.load(resp)
    
def _ingest_artifact(
    artifact_id: str,
    input_type: str,
    input_value: str,
    metadata: dict,
) -> list[str]:
    """Ingest artifact to Zylon API with flexible input type and metadata."""
    body = {
        "input": {"type": input_type, "value": input_value},
        "artifact": artifact_id,
        "collection": VECTOR_COLLECTION,
        "metadata": metadata,
    }
    print(f"Ingesting file {metadata.get('file_name')}")
    return http_post(f"{API_HOST}/v1/artifacts/ingest", body, TOKEN)


def ingest_from_file(
    artifact_id: str, file_path: str, stock_name: str, token: str = TOKEN
) -> list[str]:
    """Ingest a local file to Zylon."""
    filename = file_path.split("/")[-1]
    with open(file_path, "rb") as f:
        file_bytes = f.read()
    file_base64 = base64.b64encode(file_bytes).decode("utf-8")
    metadata = {"file_name": filename, "stock_name": stock_name}
    return _ingest_artifact(
        artifact_id=artifact_id,
        input_type="file",
        input_value=file_base64,
        metadata=metadata,
    )

def ingest_from_uri(
    artifact_id: str, file_uri: str, stock_name: str, token: str = TOKEN
) -> list[str]:
    """Ingest a file from URI to Zylon."""
    filename = file_uri.split("/")[-1]
    metadata = {"file_name": filename, "stock_name": stock_name}
    return _ingest_artifact(
        artifact_id=artifact_id,
        input_type="uri",
        input_value=file_uri,
        metadata=metadata,
    )

def send_chat_message(
    query: str,
    collection_name: Optional[str] = None,
    artifact_ids: Optional[List[str]] = None,
    tools: Optional[List[Dict]] = None,
) -> Dict:
    """
    Send a chat message to Zylon using tools.
    """
    body = {
            "messages": [{"role": "user", "content": query}],
        }
    if artifact_ids and collection_name:
        tool_context = [
            {   'type': 'ingested_artifact',
                'context_filter':{
                'collection': collection_name,
                'artifacts': artifact_ids
            }}]
        body["tool_context"] = tool_context
    if tools:
        body["tools"] = tools

    return http_post(f"{API_HOST}/v1/messages", body, TOKEN)
    

def generate_price_csv(ticker: str, months: int = 12) -> str:
    """
    Generate a CSV file with historical price data for a given stock ticker.

    The CSV file is saved locally and the path to the file is returned.
    """
    end_date = datetime.now()
    start_date = end_date - timedelta(days=30 * months)
    df = yf.download(ticker, start=start_date.strftime("%Y-%m-%d"), end=end_date.strftime("%Y-%m-%d"), auto_adjust=True)

    if df.empty:
        raise Exception(f"No data found for ticker {ticker}")

    csv_path = f"./{ticker}_price.csv"
    df.to_csv(csv_path)
    print(f"Price data for {ticker} saved to {csv_path}")
    return csv_path

def remove_price_csv(ticker: str):
    """
    Remove the local CSV file for a given stock ticker.
    """
    csv_path = f"./{ticker}_price.csv"
    try:
        os.remove(csv_path)
        print(f"Removed local price CSV file: {csv_path}")
    except FileNotFoundError:
        print(f"No local price CSV file found to remove: {csv_path}")


def format_response(query: str, response: Dict) -> str:
    """
    Extract and format the assistant's response from the Zylon API response.
    """
    return f"""
    Query: {query}
    Response: {response["content"][0]["text"]}
    """

def extract_response_text(response: dict) -> str:
    """
    Extract the main text response from the Zylon API response.
    """
    contents = response.get("content", [])
    for item in reversed(contents):  # Extract content from the end
        if isinstance(item, dict) and "text" in item:
            return item["text"]
    return ""

def format_tool_response(query: str, tool_response: Dict) -> str:
    """
    Extract and format the tool's response from the Zylon API response.
    """
    return f"""
    Query: {query}
    Tool Response: {extract_response_text(tool_response)}
    """

if __name__ == '__main__':
    ticker = "AAPL"
    example_files = [
        "https://www.apple.com/newsroom/pdfs/fy2025-q2/FY25_Q2_Consolidated_Financial_Statements.pdf",
        "https://www.apple.com/newsroom/pdfs/fy2025-q3/FY25_Q3_Consolidated_Financial_Statements.pdf"
    ]
    
    # Ingest uris documents
    documents_artifact_ids = []
    for file_url in example_files:
        artifact_id = str(uuid.uuid4())
        documents_artifact_ids.append(artifact_id)
        ingest_from_uri(artifact_id, file_url, ticker)

    # Generate, ingest and remove price CSV
    price_csv_path = generate_price_csv(ticker)
    price_artifact_id = str(uuid.uuid4())
    ingest_from_file(price_artifact_id, price_csv_path, ticker)
    remove_price_csv(ticker)

    # Query documents
    questions = [
        "Total net sales in march 2025 and march 2024?",
    ]

    responses = []
    for question in questions:
        response = send_chat_message(
        artifact_ids=documents_artifact_ids,
        collection_name=VECTOR_COLLECTION,
        query=question,
        tools=[
            {"name": "semantic_search", "type": "semantic_search_v1"},
        ])
        responses.append(format_tool_response(question, response))

    print("========= Document Responses: =========")
    for r in responses:
        print(r)


    # Query price data
    questions_price = [
        "What is the day with the highest closing price?",
        "What is the day with the lowest closing price? Include only values not nan",
    ]
    responses_price = []
    for question in questions_price:
        response = send_chat_message(
        artifact_ids=[price_artifact_id],
        collection_name=VECTOR_COLLECTION,
        query=question,
        tools=[
            {"name": "tabular_analysis", "type": "tabular_analysis_v1"},
        ])
        responses_price.append(format_tool_response(question, response))

    print("========= Price Data Responses: =========")
    for r in responses_price:
        print(r)

    ## Final query with all contexts
    all_responses = responses + responses_price
    prompt = """
    Given the following information extracted from company documents, earnings reports, and historical stock price data,
    ---
    {responses}
    ---
    Provide a concise summary of the key financial highlights and stock performance.
    """.format(responses="\n".join(all_responses))
    final_answer = send_chat_message(
        query=prompt,
    )

    print("========= Final Answer: =========")
    print(extract_response_text(final_answer))