Update app.py

app.py

@@ -120,45 +120,35 @@ except Exception as e:
 
 # try_claude_api 함수 수정
 async def try_claude_api(system_message, claude_messages, timeout=15):
-    if not claude_client:
-        raise Exception("Claude client not initialized")
-    
     try:
         start_time = time.time()
-        
-        # asyncio.timeout() 대신 asyncio.wait_for() 사용
-        async def stream_content():
-            with claude_client.messages.stream(
-                model="claude-3-5-sonnet-20241022",
-                max_tokens=7800,
-                system=system_message,
-                messages=claude_messages
-            ) as stream:
-                collected_content = ""
-                async for chunk in stream:
-                    if chunk.type == "content_block_delta":
-                        collected_content += chunk.delta.text
-                        yield collected_content
-                        await asyncio.sleep(0)
+        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:
+                current_time = time.time()
+                if current_time - start_time > timeout:
+                    print(f"Claude API response time: {current_time - start_time:.2f} seconds")
+                    raise TimeoutError("Claude API timeout")
+                if chunk.type == "content_block_delta":
+                    collected_content += chunk.delta.text
+                    yield collected_content
+                    await asyncio.sleep(0)
+                
+                start_time = current_time
                 
-        # asyncio.wait_for()로 타임아웃 설정
-        async for content in asyncio.wait_for(stream_content(), timeout):
-            yield content
-            
-    except asyncio.TimeoutError:
-        raise TimeoutError(f"Claude API timeout after {timeout} seconds")
     except Exception as e:
         print(f"Claude API error: {str(e)}")
-        raise
+        raise e
 
-# try_openai_api 함수 수정 
 async def try_openai_api(openai_messages):
-    if not openai_client:
-        raise Exception("OpenAI client not initialized")
-    
     try:
-        response = await openai_client.chat.completions.create(
-            model="gpt-4",
+        stream = openai_client.chat.completions.create(
+            model="gpt-4o",
             messages=openai_messages,
             stream=True,
             max_tokens=4096,
@@ -166,14 +156,14 @@ async def try_openai_api(openai_messages):
         )
         
         collected_content = ""
-        async for chunk in response:
-            if chunk.choices[0].delta.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
+        raise e
 
 class Demo:
     def __init__(self):
@@ -186,24 +176,24 @@ class Demo:
         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() != ''
+        ]
+        
+        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:
-            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() != ''
-            ]
-            
-            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})
-            
             yield [
                 "Generating code...",
                 _history,
@@ -211,14 +201,11 @@ class Demo:
                 gr.update(active_key="loading"),
                 gr.update(open=True)
             ]
-            
+            await asyncio.sleep(0)
+
             collected_content = None
-            error_message = None
-            
-            # Claude API 시도
             try:
                 async for content in try_claude_api(system_message, claude_messages):
-                    collected_content = content
                     yield [
                         content,
                         _history,
@@ -226,25 +213,23 @@ class Demo:
                         gr.update(active_key="loading"),
                         gr.update(open=True)
                     ]
+                    await asyncio.sleep(0)
+                collected_content = content
+                
             except Exception as claude_error:
-                print(f"Claude API error: {str(claude_error)}")
-                error_message = str(claude_error)
+                print(f"Falling back to OpenAI API due to Claude error: {str(claude_error)}")
                 
-                # OpenAI API 시도
-                try:
-                    async for content in try_openai_api(openai_messages):
-                        collected_content = content
-                        yield [
-                            content,
-                            _history,
-                            None,
-                            gr.update(active_key="loading"),
-                            gr.update(open=True)
-                        ]
-                except Exception as openai_error:
-                    print(f"OpenAI API error: {str(openai_error)}")
-                    error_message = f"Both APIs failed: Claude - {str(claude_error)}, OpenAI - {str(openai_error)}"
-            
+                async for content in try_openai_api(openai_messages):
+                    yield [
+                        content,
+                        _history,
+                        None,
+                        gr.update(active_key="loading"),
+                        gr.update(open=True)
+                    ]
+                    await asyncio.sleep(0)
+                collected_content = content
+
             if collected_content:
                 _history = messages_to_history([
                     {'role': Role.SYSTEM, 'content': system_message}
@@ -261,21 +246,16 @@ class Demo:
                     gr.update(open=True)
                 ]
             else:
-                raise ValueError(error_message or "No content was generated")
+                raise ValueError("No content was generated from either API")
                 
         except Exception as e:
             print(f"Error details: {str(e)}")
-            yield [
-                f"Error: {str(e)}",
-                _history,
-                None,
-                gr.update(active_key="empty"),
-                gr.update(open=False)
-            ]
+            raise ValueError(f'Error calling APIs: {str(e)}')
 
     def clear_history(self):
         return []
 
+
 def remove_code_block(text):
     pattern = r'```html\n(.+?)\n```'
     match = re.search(pattern, text, re.DOTALL)