final version v1

api.py

@@ -5,6 +5,7 @@ import traceback
 from typing import AsyncGenerator, List, Dict
 from config import logger
 
+# ===== OpenAI =====
 async def ask_openai(query: str, history: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
     openai_api_key = os.getenv("OPENAI_API_KEY")
     if not openai_api_key:
@@ -12,12 +13,11 @@ async def ask_openai(query: str, history: List[Dict[str, str]]) -> AsyncGenerato
         yield "Error: OpenAI API key not provided."
         return
 
-    # Build message history with user and assistant roles
     messages = []
     for msg in history:
         messages.append({"role": "user", "content": msg["user"]})
-        if msg["bot"]:
-            messages.append({"role": "assistant", "content": msg["bot"]})
+        if msg.get("openai"):
+            messages.append({"role": "assistant", "content": msg["openai"]})
     messages.append({"role": "user", "content": query})
 
     headers = {
@@ -43,32 +43,24 @@ async def ask_openai(query: str, history: List[Dict[str, str]]) -> AsyncGenerato
                             line, buffer = buffer.split("\n", 1)
                             if line.startswith("data: "):
                                 data = line[6:]
-                                if data == "[DONE]":
+                                if data.strip() == "[DONE]":
                                     break
                                 if not data.strip():
                                     continue
                                 try:
                                     json_data = json.loads(data)
-                                    if "choices" in json_data and json_data["choices"]:
-                                        delta = json_data["choices"][0].get("delta", {})
-                                        if "content" in delta and delta["content"] is not None:
-                                            logger.info(f"OpenAI yielding chunk: {delta['content']}")
-                                            yield delta["content"]
-                                except json.JSONDecodeError as e:
-                                    logger.error(f"Error parsing OpenAI stream chunk: {str(e)} - Data: {data}")
-                                    yield f"Error parsing stream: {str(e)}"
+                                    delta = json_data["choices"][0].get("delta", {})
+                                    if "content" in delta:
+                                        yield delta["content"]
                                 except Exception as e:
-                                    logger.error(f"Unexpected error in OpenAI stream: {str(e)} - Data: {data}")
-                                    yield f"Error in stream: {str(e)}"
-
-    except httpx.HTTPStatusError as e:
-        response_text = await e.response.aread()
-        logger.error(f"OpenAI HTTP Status Error: {e.response.status_code}, {response_text}")
-        yield f"Error: OpenAI HTTP Status Error: {e.response.status_code}, {response_text.decode('utf-8')}"
+                                    logger.error(f"OpenAI parse error: {e}")
+                                    yield f"[OpenAI Error]: {e}"
     except Exception as e:
-        logger.error(f"OpenAI Error: {str(e)}")
-        yield f"Error: OpenAI Error: {str(e)}"
+        logger.error(f"OpenAI API error: {e}")
+        yield f"[OpenAI Error]: {e}"
 
+
+# ===== Anthropic =====
 async def ask_anthropic(query: str, history: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
     anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
     if not anthropic_api_key:
@@ -79,8 +71,8 @@ async def ask_anthropic(query: str, history: List[Dict[str, str]]) -> AsyncGener
     messages = []
     for msg in history:
         messages.append({"role": "user", "content": msg["user"]})
-        if msg["bot"]:
-            messages.append({"role": "assistant", "content": msg["bot"]})
+        if msg.get("anthropic"):
+            messages.append({"role": "assistant", "content": msg["anthropic"]})
     messages.append({"role": "user", "content": query})
 
     headers = {
@@ -97,8 +89,7 @@ async def ask_anthropic(query: str, history: List[Dict[str, str]]) -> AsyncGener
     }
 
     try:
-        async with httpx.AsyncClient(timeout=30.0) as client:
-            logger.info(f"Sending Anthropic streaming request: {payload}")
+        async with httpx.AsyncClient() as client:
             async with client.stream("POST", "https://api.anthropic.com/v1/messages", headers=headers, json=payload) as response:
                 response.raise_for_status()
                 buffer = ""
@@ -115,26 +106,17 @@ async def ask_anthropic(query: str, history: List[Dict[str, str]]) -> AsyncGener
                                     continue
                                 try:
                                     json_data = json.loads(data)
-                                    if json_data.get("type") == "content_block_delta" and "delta" in json_data and "text" in json_data["delta"]:
-                                        logger.info(f"Anthropic yielding chunk: {json_data['delta']['text']}")
-                                        yield json_data["delta"]["text"]
-                                    elif json_data.get("type") == "message_start" or json_data.get("type") == "message_delta":
-                                        continue
-                                except json.JSONDecodeError as e:
-                                    logger.error(f"Error parsing Anthropic stream chunk: {str(e)} - Data: {data}")
-                                    yield f"Error parsing stream: {str(e)}"
+                                    if json_data.get("type") == "content_block_delta" and "delta" in json_data:
+                                        yield json_data["delta"].get("text", "")
                                 except Exception as e:
-                                    logger.error(f"Unexpected error in Anthropic stream: {str(e)} - Data: {data}")
-                                    yield f"Error in stream: {str(e)}"
-
-    except httpx.HTTPStatusError as e:
-        response_text = await e.response.aread()
-        logger.error(f"Anthropic HTTP Status Error: {e.response.status_code}, {response_text.decode('utf-8')}")
-        yield f"Error: Anthropic HTTP Status Error: {e.response.status_code}, {response_text.decode('utf-8')}"
+                                    logger.error(f"Anthropic parse error: {e}")
+                                    yield f"[Anthropic Error]: {e}"
     except Exception as e:
-        logger.error(f"Anthropic Error: {str(e)}\nStack trace: {traceback.format_exc()}")
-        yield f"Error: Anthropic Error: {str(e)}"
+        logger.error(f"Anthropic API error: {e}")
+        yield f"[Anthropic Error]: {e}"
+
 
+# ===== Gemini =====
 async def ask_gemini(query: str, history: List[Dict[str, str]]) -> AsyncGenerator[str, None]:
     gemini_api_key = os.getenv("GEMINI_API_KEY")
     if not gemini_api_key:
@@ -144,20 +126,18 @@ async def ask_gemini(query: str, history: List[Dict[str, str]]) -> AsyncGenerato
 
     history_text = ""
     for msg in history:
-        history_text += f"User: {msg['user']}\nAssistant: {msg['bot']}\n" if msg['bot'] else f"User: {msg['user']}\n"
-    full_query = history_text + f"User: {query}\n"
-
-    headers = {
-        "Content-Type": "application/json"
-    }
+        history_text += f"User: {msg['user']}\n"
+        if msg.get("gemini"):
+            history_text += f"Assistant: {msg['gemini']}\n"
+    full_prompt = f"{history_text}User: {query}\n"
 
+    headers = {"Content-Type": "application/json"}
     payload = {
-        "contents": [{"parts": [{"text": full_query}]}]
+        "contents": [{"parts": [{"text": full_prompt}]}]
     }
 
     try:
-        async with httpx.AsyncClient(timeout=30.0) as client:
-            logger.info(f"Sending Gemini streaming request: {payload}")
+        async with httpx.AsyncClient() as client:
             async with client.stream(
                 "POST",
                 f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent?key={gemini_api_key}",
@@ -167,63 +147,25 @@ async def ask_gemini(query: str, history: List[Dict[str, str]]) -> AsyncGenerato
                 response.raise_for_status()
                 buffer = ""
                 async for chunk in response.aiter_text():
-                    if chunk:
-                        buffer += chunk
-                        logger.info(f"Gemini stream chunk: {chunk}")
-                        while buffer.strip():
-                            try:
-                                json_data = json.loads(buffer)
-                                logger.info(f"Parsed Gemini JSON: {json_data}")
-                                buffer = ""
-                                objects = json_data if isinstance(json_data, list) else [json_data]
-                                for obj in objects:
-                                    if isinstance(obj, dict) and "candidates" in obj and obj["candidates"]:
-                                        content = obj["candidates"][0].get("content", {})
-                                        if "parts" in content and content["parts"]:
-                                            text = content["parts"][0].get("text", "")
-                                            if text:
-                                                logger.info(f"Gemini yielding chunk: {text}")
-                                                yield text
-                                break
-                            except json.JSONDecodeError:
-                                brace_count = 0
-                                split_index = -1
-                                for i, char in enumerate(buffer):
-                                    if char == '{':
-                                        brace_count += 1
-                                    elif char == '}':
-                                        brace_count -= 1
-                                        if brace_count == 0:
-                                            split_index = i + 1
-                                if split_index > 0:
-                                    try:
-                                        json_str = buffer[:split_index]
-                                        json_data = json.loads(json_str)
-                                        logger.info(f"Parsed Gemini JSON: {json_data}")
-                                        buffer = buffer[split_index:].lstrip(',')
-                                        objects = json_data if isinstance(json_data, list) else [json_data]
-                                        for obj in objects:
-                                            if isinstance(obj, dict) and "candidates" in obj and obj["candidates"]:
-                                                content = obj["candidates"][0].get("content", {})
-                                                if "parts" in content and content["parts"]:
-                                                    text = content["parts"][0].get("text", "")
-                                                    if text:
-                                                        logger.info(f"Gemini yielding chunk: {text}")
-                                                        yield text
-                                        continue
-                                    except json.JSONDecodeError:
-                                        pass
-                                break
-                            except Exception as e:
-                                logger.error(f"Unexpected error in Gemini stream: {str(e)} - Buffer: {buffer}")
-                                yield f"Error in stream: {str(e)}"
-                                buffer = ""
-                                break
-
-    except httpx.HTTPStatusError as e:
-        response_text = await e.response.aread()
-        logger.error(f"Gemini HTTP Status Error: {e.response.status_code}, {response_text.decode('utf-8')}")
-        yield f"Error: Gemini HTTP Status Error: {e.response.status_code}, {response_text.decode('utf-8')}"
+                    if not chunk.strip():
+                        continue
+                    buffer += chunk
+                    try:
+                        json_data = json.loads(buffer.strip(", \n"))
+                        buffer = ""
+
+                        # handle both list and dict format
+                        objects = json_data if isinstance(json_data, list) else [json_data]
+                        for obj in objects:
+                            candidates = obj.get("candidates", [])
+                            if candidates:
+                                parts = candidates[0].get("content", {}).get("parts", [])
+                                for part in parts:
+                                    text = part.get("text", "")
+                                    if text:
+                                        yield text
+                    except json.JSONDecodeError:
+                        continue  # wait for more data
     except Exception as e:
-        logger.error(f"Gemini Error: {str(e)}\nStack trace: {traceback.format_exc()}")
-        yield f"Error: Gemini Error: {str(e)}"
+        logger.error(f"Gemini API error: {e}")
+        yield f"[Gemini Error]: {e}"

app.py

@@ -5,32 +5,33 @@ from core import submit_query, clear_history
 with gr.Blocks(theme=gr.themes.Soft(), css=".full-height { height: 100%; display: flex; align-items: stretch; min-height: 40px; } .full-height button { height: 100%; padding: 8px 16px; } .providers-row { height: 100%; display: flex; align-items: stretch; min-height: 40px; } .providers-row .checkbox-group { height: 100%; display: flex; flex-direction: row; align-items: center; gap: 10px; }") as demo:
     gr.Markdown("# Multi-Model Chat")
     gr.Markdown("Chat with OpenAI, Anthropic, or Gemini. Select providers and compare responses side by side!")
-    
+
     with gr.Row(elem_classes="providers-row"):
         providers = gr.CheckboxGroup(choices=["OpenAI", "Anthropic", "Gemini"], label="Select Providers", value=["OpenAI"], elem_classes="checkbox-group")
-    
+
     with gr.Row(elem_classes="full-height"):
         query = gr.Textbox(label="Enter your query", placeholder="e.g., What is the capital of the United States?", scale=4)
         submit_button = gr.Button("Submit", scale=1)
-    
+
     with gr.Row():
         clear_button = gr.Button("Clear History")
-    
+
     with gr.Row():
         openai_chatbot = gr.Chatbot(label="OpenAI", type="messages", scale=1)
         anthropic_chatbot = gr.Chatbot(label="Anthropic", type="messages", scale=1)
         gemini_chatbot = gr.Chatbot(label="Gemini", type="messages", scale=1)
 
+    chat_history = gr.State([])
+
     submit_button.click(
         fn=submit_query,
-        inputs=[query, providers],
-        outputs=[query, openai_chatbot, anthropic_chatbot, gemini_chatbot]
+        inputs=[query, providers, chat_history],
+        outputs=[query, openai_chatbot, anthropic_chatbot, gemini_chatbot, chat_history]
     )
     clear_button.click(
         fn=clear_history,
         inputs=[],
-        outputs=[openai_chatbot, anthropic_chatbot, gemini_chatbot]
+        outputs=[openai_chatbot, anthropic_chatbot, gemini_chatbot, chat_history]
     )
 
-# Launch the Gradio app
-demo.launch()
+demo.launch()

core.py

@@ -3,30 +3,37 @@ from typing import AsyncGenerator, List, Dict, Tuple
 from config import logger
 from api import ask_openai, ask_anthropic, ask_gemini
 
-async def query_model(query: str, providers: List[str], history: List[Dict[str, str]]) -> AsyncGenerator[Tuple[str, List[Dict[str, str]], List[Dict[str, str]], List[Dict[str, str]]], None]:
+async def query_model(
+    query: str,
+    providers: List[str],
+    history: List[Dict[str, str]]
+) -> AsyncGenerator[
+    Tuple[str, List[Dict[str, str]], List[Dict[str, str]], List[Dict[str, str]], List[Dict[str, str]]],
+    None
+]:
     logger.info(f"Processing query with providers: {providers}")
     openai_response = ""
     anthropic_response = ""
     gemini_response = ""
-    
+
     openai_messages = []
     anthropic_messages = []
     gemini_messages = []
-    
+
+    # Build message history for each provider
     for msg in history:
-        openai_messages.append({"role": "user", "content": msg["user"]})
-        anthropic_messages.append({"role": "user", "content": msg["user"]})
-        gemini_messages.append({"role": "user", "content": msg["user"]})
-        if msg["bot"]:
-            lines = msg["bot"].split("\n\n")
-            for line in lines:
-                if line.startswith("[OpenAI]:"):
-                    openai_messages.append({"role": "assistant", "content": line[len("[OpenAI]:"):].strip()})
-                elif line.startswith("[Anthropic]:"):
-                    anthropic_messages.append({"role": "assistant", "content": line[len("[Anthropic]:"):].strip()})
-                elif line.startswith("[Gemini]:"):
-                    gemini_messages.append({"role": "assistant", "content": line[len("[Gemini]:"):].strip()})
+        if "user" in msg:
+            openai_messages.append({"role": "user", "content": msg["user"]})
+            anthropic_messages.append({"role": "user", "content": msg["user"]})
+            gemini_messages.append({"role": "user", "content": msg["user"]})
+        if msg.get("openai"):
+            openai_messages.append({"role": "assistant", "content": msg["openai"]})
+        if msg.get("anthropic"):
+            anthropic_messages.append({"role": "assistant", "content": msg["anthropic"]})
+        if msg.get("gemini"):
+            gemini_messages.append({"role": "assistant", "content": msg["gemini"]})
 
+    # Append the user query and prepare for assistant response
     if "OpenAI" in providers:
         openai_messages.append({"role": "user", "content": query})
         openai_messages.append({"role": "assistant", "content": ""})
@@ -37,6 +44,10 @@ async def query_model(query: str, providers: List[str], history: List[Dict[str,
         gemini_messages.append({"role": "user", "content": query})
         gemini_messages.append({"role": "assistant", "content": ""})
 
+    # Yield initial state with user query
+    logger.info(f"Yielding initial state with user query: {query}")
+    yield "", openai_messages, anthropic_messages, gemini_messages, history
+
     tasks = []
     if "OpenAI" in providers:
         tasks.append(("OpenAI", ask_openai(query, history), openai_response, openai_messages))
@@ -45,7 +56,12 @@ async def query_model(query: str, providers: List[str], history: List[Dict[str,
     if "Gemini" in providers:
         tasks.append(("Gemini", ask_gemini(query, history), gemini_response, gemini_messages))
 
-    async def collect_chunks(provider: str, generator: AsyncGenerator, response: str, messages: List[Dict[str, str]]) -> AsyncGenerator[Tuple[str, str, List[Dict[str, str]]], None]:
+    async def collect_chunks(
+        provider: str,
+        generator: AsyncGenerator[str, None],
+        response: str,
+        messages: List[Dict[str, str]]
+    ) -> AsyncGenerator[Tuple[str, str, List[Dict[str, str]]], None]:
         async for chunk in generator:
             response += chunk
             messages[-1] = {"role": "assistant", "content": response}
@@ -57,7 +73,7 @@ async def query_model(query: str, providers: List[str], history: List[Dict[str,
     while active_generators:
         tasks_to_wait = []
         new_generator_states = []
-        
+
         for provider, gen, active_task in active_generators:
             if active_task is None or active_task.done():
                 try:
@@ -71,12 +87,12 @@ async def query_model(query: str, providers: List[str], history: List[Dict[str,
             else:
                 new_generator_states.append((provider, gen, active_task))
                 tasks_to_wait.append(active_task)
-        
+
         if not tasks_to_wait:
             break
 
         done, _ = await asyncio.wait(tasks_to_wait, return_when=asyncio.FIRST_COMPLETED)
-        
+
         for provider, gen, task in new_generator_states:
             if task in done:
                 try:
@@ -91,7 +107,7 @@ async def query_model(query: str, providers: List[str], history: List[Dict[str,
                         gemini_response = response
                         gemini_messages = messages
                     logger.info(f"Yielding update for {provider}: {response[:50]}...")
-                    yield "", openai_messages, anthropic_messages, gemini_messages
+                    yield "", openai_messages, anthropic_messages, gemini_messages, history
                     new_generator_states[new_generator_states.index((provider, gen, task))] = (provider, gen, None)
                 except StopAsyncIteration:
                     logger.info(f"Generator for {provider} completed")
@@ -102,46 +118,40 @@ async def query_model(query: str, providers: List[str], history: List[Dict[str,
 
         active_generators = new_generator_states
 
-    responses = []
-    if openai_response.strip() and not openai_response.startswith("Error:"):
-        responses.append(f"[OpenAI]: {openai_response}")
-    if anthropic_response.strip() and not anthropic_response.startswith("Error:"):
-        responses.append(f"[Anthropic]: {anthropic_response}")
-    if gemini_response.strip() and not gemini_response.startswith("Error:"):
-        responses.append(f"[Gemini]: {gemini_response}")
-    
-    combined_response = "\n\n".join(responses) if responses else "No valid responses received."
-    updated_history = history + [{"user": query, "bot": combined_response}]
-    logger.info(f"Updated history: {updated_history}")
+    updated_history = history + [{
+        "user": query,
+        "openai": openai_response.strip() if openai_response else "",
+        "anthropic": anthropic_response.strip() if anthropic_response else "",
+        "gemini": gemini_response.strip() if gemini_response else ""
+    }]
 
-    yield "", openai_messages, anthropic_messages, gemini_messages
+    logger.info(f"Updated history: {updated_history}")
+    yield "", openai_messages, anthropic_messages, gemini_messages, updated_history
 
-async def submit_query(query: str, providers: List[str]) -> AsyncGenerator[Tuple[str, List[Dict[str, str]], List[Dict[str, str]], List[Dict[str, str]]], None]:
+async def submit_query(
+    query: str,
+    providers: List[str],
+    history: List[Dict[str, str]]
+) -> AsyncGenerator[
+    Tuple[str, List[Dict[str, str]], List[Dict[str, str]], List[Dict[str, str]], List[Dict[str, str]]],
+    None
+]:
     if not query.strip():
-        openai_messages = [{"role": "assistant", "content": "Please enter a query."}]
-        anthropic_messages = [{"role": "assistant", "content": "Please enter a query."}]
-        gemini_messages = [{"role": "assistant", "content": "Please enter a query."}]
-        logger.info("Yielding empty query response")
-        yield "", openai_messages, anthropic_messages, gemini_messages
+        msg = {"role": "assistant", "content": "Please enter a query."}
+        yield "", [msg], [msg], [msg], history
         return
-    
+
     if not providers:
-        openai_messages = [{"role": "assistant", "content": "Please select at least one provider."}]
-        anthropic_messages = [{"role": "assistant", "content": "Please select at least one provider."}]
-        gemini_messages = [{"role": "assistant", "content": "Please select at least one provider."}]
-        logger.info("Yielding no providers response")
-        yield "", openai_messages, anthropic_messages, gemini_messages
+        msg = {"role": "assistant", "content": "Please select at least one provider."}
+        yield "", [msg], [msg], [msg], history
         return
 
-    history = []
-    
-    async for response_chunk, openai_messages, anthropic_messages, gemini_messages in query_model(query, providers, history):
-        logger.info(f"Submitting update to UI: OpenAI: {openai_messages[-1]['content'][:50] if openai_messages else ''}, "
-                    f"Anthropic: {anthropic_messages[-1]['content'][:50] if anthropic_messages else ''}, "
-                    f"Gemini: {gemini_messages[-1]['content'][:50] if gemini_messages else ''}")
-        yield "", openai_messages, anthropic_messages, gemini_messages
-    logger.info("Final UI update")
-    yield "", openai_messages, anthropic_messages, gemini_messages
+    async for _, openai_msgs, anthropic_msgs, gemini_msgs, updated_history in query_model(query, providers, history):
+        logger.info(f"Submitting update to UI: OpenAI: {openai_msgs[-1]['content'][:50] if openai_msgs else ''}, "
+                    f"Anthropic: {anthropic_msgs[-1]['content'][:50] if anthropic_msgs else ''}, "
+                    f"Gemini: {gemini_msgs[-1]['content'][:50] if gemini_msgs else ''}")
+        yield "", openai_msgs, anthropic_msgs, gemini_msgs, updated_history
 
 def clear_history():
-    return [], [], []
+    logger.info("Clearing history")
+    return [], [], [], []