File size: 4,512 Bytes
ec9f5b2
 
 
 
 
 
 
 
 
 
 
 
 
2c703f2
 
 
 
 
 
ec9f5b2
ad0b759
ec9f5b2
db0a93c
ec9f5b2
 
 
 
 
aacc297
 
 
 
 
 
 
 
 
 
 
ad0b759
 
 
aacc297
 
 
 
 
 
 
ad0b759
 
 
 
ec9f5b2
2c703f2
ec9f5b2
2c703f2
ec9f5b2
 
2c703f2
ad0b759
ec9f5b2
 
 
 
 
 
 
ad0b759
ec9f5b2
ad0b759
 
ec9f5b2
 
 
 
 
2c703f2
 
 
aacc297
 
ad0b759
ec9f5b2
 
2c703f2
 
ad0b759
 
2c703f2
 
ec9f5b2
ee5c938
 
 
 
ec9f5b2
2c703f2
ec9f5b2
2c703f2
ec9f5b2
 
2c703f2
ad0b759
ec9f5b2
 
 
 
 
 
 
ad0b759
ec9f5b2
ad0b759
 
ec9f5b2
 
 
 
 
2c703f2
 
 
aacc297
 
ec9f5b2
 
 
2c703f2
 
ad0b759
2c703f2
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
import os
import uuid
from langchain_openai import ChatOpenAI
import pandas as pd
from pandasai import SmartDataframe
from csv_service import clean_data
from dotenv import load_dotenv
from util_service import handle_out_of_range_float, process_answer

load_dotenv()
openai_api_keys = os.getenv("OPENAI_API_KEYS").split(",")
openai_api_base = os.getenv("OPENAI_API_BASE")

# Create an array of ChatOpenAI instances
llm_instances = [
    ChatOpenAI(model='gpt-4o', api_key=key, base_url=openai_api_base)
    for key in openai_api_keys
]
current_llm_index = 0

instructions = """
- Please ensure that each value is clearly visible, You may need to adjust the font size, rotate the labels, or use truncation to improve readability (if needed).
- For multiple charts, arrange them in a grid format (2x2, 3x3, etc.)
- Use professional and color-blind friendly palettes.
- Do not use sns.set_palette()
- Read above instructions and follow them.
"""

def should_rotate_key(response):
    """Check if response indicates API key needs rotation"""
    if isinstance(response, str):
        return any(msg in response for msg in [
            "plan.rule:api_request",
            "resource limit",
            "exhausted the available",
            "update your payment method"
        ])
    return False

def handle_api_error(error, current_index):
    """Determine if we should rotate API keys based on error message"""
    error_msg = str(error)
    if any(msg in error_msg for msg in [
        "plan.rule:api_request",
        "429",
        "resource limit",
        "exhausted the available",
        "update your payment method"
    ]):
        print(f"Rotating API key due to resource limit (key index {current_index})")
        return True
    return False

def openai_chat(csv_url: str, question: str):
    global current_llm_index

    while current_llm_index < len(llm_instances):
        try:
            data = clean_data(csv_url)
            llm = llm_instances[current_llm_index]
            
            chart_filename = f"chart_{uuid.uuid4()}.png"
            chart_path = os.path.join("generated_charts", chart_filename)
            
            df = SmartDataframe(
                data,
                config={
                    'llm': llm,
                    'save_charts': True,
                    'open_charts': False,
                    'save_charts_path': os.path.dirname(chart_path),
                    'custom_chart_filename': chart_filename
                }
            )
            
            answer = df.chat(question)
            
            # Check if answer contains API limit error or process_answer is true
            if should_rotate_key(answer) or process_answer(answer):
                current_llm_index += 1
                continue
                
            return answer

        except Exception as e:
            if handle_api_error(e, current_llm_index):
                current_llm_index += 1
                continue
            return {"error": f"Non-recoverable error: {str(e)}"}
    
    return {"error": "All API keys exhausted. Please update billing information."}





def openai_chart(csv_url: str, question: str):
    global current_llm_index

    while current_llm_index < len(llm_instances):
        try:
            data = clean_data(csv_url)
            llm = llm_instances[current_llm_index]
            
            chart_filename = f"chart_{uuid.uuid4()}.png"
            chart_path = os.path.join("generated_charts", chart_filename)
            
            df = SmartDataframe(
                data,
                config={
                    'llm': llm,
                    'save_charts': True,
                    'open_charts': False,
                    'save_charts_path': os.path.dirname(chart_path),
                    'custom_chart_filename': chart_filename
                }
            )
            
            answer = df.chat(question + instructions)
            
            # Check if answer contains API limit error or process_answer is true
            if should_rotate_key(answer) or process_answer(answer):
                current_llm_index += 1
                continue
                
            return answer

        except Exception as e:
            if handle_api_error(e, current_llm_index):
                current_llm_index += 1
                continue
            return {"error": f"Chart generation failed: {str(e)}"}
    
    return {"error": "All API keys exhausted. Please update billing information."}