Update app.py

app.py

@@ -5,6 +5,9 @@ from http import HTTPStatus
 from typing import Dict, List, Optional, Tuple
 import base64
 import anthropic
+import openai
+import asyncio
+import time
 from functools import partial
 
 import gradio as gr
@@ -54,8 +57,141 @@ def get_image_base64(image_path):
         encoded_string = base64.b64encode(image_file.read()).decode()
     return encoded_string
     
-YOUR_API_TOKEN = os.getenv('ANTHROPIC_API_KEY')
-client = anthropic.Anthropic(api_key=YOUR_API_TOKEN)
+# API 클라이언트 초기화
+YOUR_ANTHROPIC_TOKEN = os.getenv('ANTHROPIC_API_KEY')
+YOUR_OPENAI_TOKEN = os.getenv('OPENAI_API_KEY')
+
+claude_client = anthropic.Anthropic(api_key=YOUR_ANTHROPIC_TOKEN)
+openai_client = openai.OpenAI(api_key=YOUR_OPENAI_TOKEN)
+
+
+
+async def try_claude_api(system_message, claude_messages, timeout=5):
+    try:
+        start_time = time.time()
+        with claude_client.messages.stream(
+            model="claude-3-5-sonnet-20241022",
+            max_tokens=7800,
+            system=system_message,
+            messages=claude_messages
+        ) as stream:
+            collected_content = ""
+            for chunk in stream:
+                if time.time() - start_time > timeout:
+                    raise TimeoutError("Claude API timeout")
+                if chunk.type == "content_block_delta":
+                    collected_content += chunk.delta.text
+                    yield collected_content
+    except Exception as e:
+        print(f"Claude API error: {str(e)}")
+        raise e
+
+
+        
+async def try_openai_api(openai_messages):
+    try:
+        stream = openai_client.chat.completions.create(
+            model="gpt-4o",
+            messages=openai_messages,
+            stream=True,
+            max_tokens=4096,
+            temperature=0.7
+        )
+        
+        collected_content = ""
+        for chunk in stream:
+            if chunk.choices[0].delta.content is not None:
+                collected_content += chunk.choices[0].delta.content
+                yield collected_content
+                
+    except Exception as e:
+        print(f"OpenAI API error: {str(e)}")
+        raise e
+
+async def generation_code(query: Optional[str], _setting: Dict[str, str], _history: Optional[History]):
+    if not query or query.strip() == '':
+        query = random.choice(DEMO_LIST)['description']
+    
+    if _history is None:
+        _history = []
+    
+    messages = history_to_messages(_history, _setting['system'])
+    system_message = messages[0]['content']
+    
+    # Claude 메시지 포맷
+    claude_messages = [
+        {"role": msg["role"] if msg["role"] != "system" else "user", "content": msg["content"]} 
+        for msg in messages[1:] + [{'role': Role.USER, 'content': query}]
+        if msg["content"].strip() != ''
+    ]
+    
+    # OpenAI 메시지 포맷
+    openai_messages = [{"role": "system", "content": system_message}]
+    for msg in messages[1:]:
+        openai_messages.append({
+            "role": msg["role"],
+            "content": msg["content"]
+        })
+    openai_messages.append({"role": "user", "content": query})
+    
+    try:
+        yield [
+            "Generating code...",
+            _history,
+            None,
+            gr.update(active_key="loading"),
+            gr.update(open=True)
+        ]
+
+        collected_content = None
+        # Claude API 시도
+        try:
+            async for content in try_claude_api(system_message, claude_messages):
+                yield [
+                    content,
+                    _history,
+                    None,
+                    gr.update(active_key="loading"),
+                    gr.update(open=True)
+                ]
+                collected_content = content
+            
+        except Exception as claude_error:
+            print(f"Falling back to OpenAI API due to Claude error: {str(claude_error)}")
+            
+            # OpenAI API로 폴백
+            async for content in try_openai_api(openai_messages):
+                yield [
+                    content,
+                    _history,
+                    None,
+                    gr.update(active_key="loading"),
+                    gr.update(open=True)
+                ]
+                collected_content = content
+
+        if collected_content:
+            _history = messages_to_history([
+                {'role': Role.SYSTEM, 'content': system_message}
+            ] + claude_messages + [{
+                'role': Role.ASSISTANT,
+                'content': collected_content
+            }])
+            
+            yield [
+                collected_content,
+                _history,
+                send_to_sandbox(remove_code_block(collected_content)),
+                gr.update(active_key="render"),
+                gr.update(open=True)
+            ]
+        else:
+            raise ValueError("No content was generated from either API")
+    
+    except Exception as e:
+        print(f"Error details: {str(e)}")
+        raise ValueError(f'Error calling APIs: {str(e)}')
+        
 
 class Role:
     SYSTEM = "system"
@@ -108,69 +244,7 @@ with gr.Blocks(css_paths="app.css",theme=theme) as demo:
         "system": SystemPrompt,
     })
 
-    def generation_code(query: Optional[str], _setting: Dict[str, str], _history: Optional[History]):
-        if not query or query.strip() == '':
-            query = random.choice(DEMO_LIST)['description']  # 빈 쿼리일 경우 랜덤 예제 사용
-        
-        if _history is None:
-            _history = []
-        
-        messages = history_to_messages(_history, _setting['system'])
-        system_message = messages[0]['content']
-        
-        claude_messages = [
-            {"role": msg["role"] if msg["role"] != "system" else "user", "content": msg["content"]} 
-            for msg in messages[1:] + [{'role': Role.USER, 'content': query}]
-            if msg["content"].strip() != ''  # 빈 메시지 필터링
-        ]
-        
-        try:
-            yield [
-                "Generating code...",
-                _history,
-                None,
-                gr.update(active_key="loading"),
-                gr.update(open=True)
-            ]
-    
-            with client.messages.stream(
-                model="claude-3-5-sonnet-20241022",
-                max_tokens=7800,
-                system=system_message,
-                messages=claude_messages
-            ) as stream:
-                collected_content = ""
-                for chunk in stream:
-                    if chunk.type == "content_block_delta":
-                        delta = chunk.delta.text
-                        collected_content += delta
-                        
-                        yield [
-                            collected_content,
-                            _history,
-                            None,
-                            gr.update(active_key="loading"),
-                            gr.update(open=True)
-                        ]
-            
-                _history = messages_to_history([
-                    {'role': Role.SYSTEM, 'content': system_message}
-                ] + claude_messages + [{
-                    'role': Role.ASSISTANT,
-                'content': collected_content
-                }])
-            
-                yield [
-                    collected_content,
-                    _history,
-                    send_to_sandbox(remove_code_block(collected_content)),
-                    gr.update(active_key="render"),
-                    gr.update(open=True)
-                ]
-        
-        except Exception as e:
-            print(f"Error details: {str(e)}")  # 디버깅을 위한 에러 출력
-            raise ValueError(f'Error calling Claude API: {str(e)}')
+
 
     with ms.Application() as app:
         with antd.ConfigProvider():
@@ -361,11 +435,13 @@ with gr.Blocks(css_paths="app.css",theme=theme) as demo:
                 open=False), inputs=[], outputs=[history_drawer])
 
             btn.click(
-                generation_code,
+                lambda: asyncio.create_task(generation_code(input.value, setting.value, history.value)),
                 inputs=[input, setting, history],
                 outputs=[code_output, history, sandbox, state_tab, code_drawer]
             )
             
+
+            
             clear_btn.click(clear_history, inputs=[], outputs=[history])
 
 if __name__ == "__main__":