All working except google streaming

app.py

@@ -82,11 +82,12 @@ async def ask_openai(query: str, history: List[Dict[str, str]]) -> AsyncGenerato
         logger.error(f"OpenAI Error: {str(e)}")
         yield f"Error: OpenAI Error: {str(e)}"
 
-async def ask_anthropic(query: str, history: List[Dict[str, str]]) -> str:
+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:
         logger.error("Anthropic API key not provided")
-        return "Error: Anthropic API key not provided."
+        yield "Error: Anthropic API key not provided."
+        return
 
     # Build message history with user and assistant roles
     messages = []
@@ -105,42 +106,55 @@ async def ask_anthropic(query: str, history: List[Dict[str, str]]) -> str:
     payload = {
         "model": "claude-3-5-sonnet-20241022",
         "max_tokens": 1024,
-        "messages": messages
+        "messages": messages,
+        "stream": True
     }
 
     try:
         async with httpx.AsyncClient(timeout=30.0) as client:
-            logger.info(f"Sending Anthropic request: {payload}")
-            response = await client.post("https://api.anthropic.com/v1/messages", headers=headers, json=payload)
-        
-        response.raise_for_status()
-        response_json = response.json()
-        logger.info(f"Anthropic response: {response_json}")
-        
-        # Validate response structure
-        if not isinstance(response_json, dict) or "content" not in response_json or not response_json["content"]:
-            logger.error(f"Invalid Anthropic response structure: {response_json}")
-            return f"Error: Invalid Anthropic response structure"
-        
-        content = response_json["content"]
-        if not isinstance(content, list) or not content or "text" not in content[0]:
-            logger.error(f"Invalid Anthropic content format: {content}")
-            return f"Error: Invalid Anthropic content format"
-        
-        return response_json["content"][0]["text"]
+            logger.info(f"Sending Anthropic streaming request: {payload}")
+            async with client.stream("POST", "https://api.anthropic.com/v1/messages", headers=headers, json=payload) as response:
+                response.raise_for_status()
+                buffer = ""
+                async for chunk in response.aiter_text():
+                    if chunk:
+                        buffer += chunk
+                        # Process complete JSON lines
+                        while "\n" in buffer:
+                            line, buffer = buffer.split("\n", 1)
+                            if line.startswith("data: "):
+                                data = line[6:]  # Remove "data: " prefix
+                                if data.strip() == "[DONE]":
+                                    break
+                                if not data.strip():
+                                    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"]:
+                                        yield json_data["delta"]["text"]
+                                    elif json_data.get("type") == "message_start" or json_data.get("type") == "message_delta":
+                                        continue  # Skip metadata events
+                                except json.JSONDecodeError as e:
+                                    logger.error(f"Error parsing Anthropic stream chunk: {str(e)} - Data: {data}")
+                                    yield f"Error parsing stream: {str(e)}"
+                                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')}")
-        return f"Error: 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')}"
     except Exception as e:
         logger.error(f"Anthropic Error: {str(e)}\nStack trace: {traceback.format_exc()}")
-        return f"Error: Anthropic Error: {str(e)}"
+        yield f"Error: Anthropic Error: {str(e)}"
 
-async def ask_gemini(query: str, history: List[Dict[str, str]]) -> str:
+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:
         logger.error("Gemini API key not provided")
-        return "Error: Gemini API key not provided."
+        yield "Error: Gemini API key not provided."
+        return
 
     # Concatenate history as text for Gemini
     history_text = ""
@@ -158,45 +172,71 @@ async def ask_gemini(query: str, history: List[Dict[str, str]]) -> str:
 
     try:
         async with httpx.AsyncClient(timeout=30.0) as client:
-            logger.info(f"Sending Gemini request: {payload}")
-            response = await client.post(
-                f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:generateContent?key={gemini_api_key}",
+            logger.info(f"Sending Gemini streaming request: {payload}")
+            async with client.stream(
+                "POST",
+                f"https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent?key={gemini_api_key}",
                 headers=headers,
                 json=payload
-            )
-        
-        response.raise_for_status()
-        response_json = response.json()
-        logger.info(f"Gemini response: {response_json}")
-        
-        # Validate response structure
-        if not isinstance(response_json, dict) or "candidates" not in response_json or not response_json["candidates"]:
-            logger.error(f"Invalid Gemini response structure: {response_json}")
-            return f"Error: Invalid Gemini response structure"
-        
-        candidates = response_json["candidates"]
-        if not isinstance(candidates, list) or not candidates or "content" not in candidates[0]:
-            logger.error(f"Invalid Gemini candidates format: {candidates}")
-            return f"Error: Invalid Gemini candidates format"
-        
-        content = candidates[0]["content"]
-        if not isinstance(content, dict) or "parts" not in content or not content["parts"]:
-            logger.error(f"Invalid Gemini content format: {content}")
-            return f"Error: Invalid Gemini content format"
-        
-        parts = content["parts"]
-        if not isinstance(parts, list) or not parts or "text" not in parts[0]:
-            logger.error(f"Invalid Gemini parts format: {parts}")
-            return f"Error: Invalid Gemini parts format"
-        
-        return parts[0]["text"]
+            ) as response:
+                response.raise_for_status()
+                buffer = ""
+                async for chunk in response.aiter_text():
+                    if chunk:
+                        buffer += chunk
+                        logger.info(f"Gemini stream chunk: {chunk}")
+                        # Try to parse buffer as multiple JSON objects
+                        while buffer.strip():
+                            try:
+                                # Attempt to parse the buffer as JSON
+                                json_data = json.loads(buffer)
+                                logger.info(f"Parsed Gemini JSON: {json_data}")
+                                buffer = ""  # Reset buffer after successful parse
+                                # Handle both single object and list of objects
+                                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:
+                                                yield text
+                                break  # Exit loop after successful parse
+                            except json.JSONDecodeError as e:
+                                # Check if buffer might contain a partial object followed by a comma
+                                comma_index = buffer.rfind(",")
+                                if comma_index != -1:
+                                    # Try parsing up to the last comma
+                                    try:
+                                        json_data = json.loads(buffer[:comma_index])
+                                        logger.info(f"Parsed Gemini JSON (before comma): {json_data}")
+                                        buffer = buffer[comma_index + 1:].strip()
+                                        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:
+                                                        yield text
+                                        continue  # Continue processing remaining buffer
+                                    except json.JSONDecodeError:
+                                        pass  # Continue accumulating buffer
+                                # If parsing fails, accumulate more data
+                                break
+                            except Exception as e:
+                                logger.error(f"Unexpected error in Gemini stream: {str(e)} - Data: {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')}")
-        return f"Error: 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')}"
     except Exception as e:
         logger.error(f"Gemini Error: {str(e)}\nStack trace: {traceback.format_exc()}")
-        return f"Error: Gemini Error: {str(e)}"
+        yield f"Error: Gemini Error: {str(e)}"
 
 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]:
     logger.info(f"Processing query with providers: {providers}")
@@ -225,30 +265,40 @@ async def query_model(query: str, providers: List[str], history: List[Dict[str,
                 elif line.startswith("[Gemini]:"):
                     gemini_messages.append({"role": "assistant", "content": line[len("[Gemini]:"):].strip()})
 
-    # Handle OpenAI (streaming)
+    # Append the current query to all chatbots
     if "OpenAI" in providers:
         openai_messages.append({"role": "user", "content": query})
+        openai_messages.append({"role": "assistant", "content": ""})
+    if "Anthropic" in providers:
+        anthropic_messages.append({"role": "user", "content": query})
+        anthropic_messages.append({"role": "assistant", "content": ""})
+    if "Gemini" in providers:
+        gemini_messages.append({"role": "user", "content": query})
+        gemini_messages.append({"role": "assistant", "content": ""})
+
+    # Handle OpenAI (streaming)
+    if "OpenAI" in providers:
         async for chunk in ask_openai(query, history):
             openai_response += chunk
-            # Yield streaming updates for OpenAI
+            # Update OpenAI chatbot with streaming response
             openai_messages[-1] = {"role": "assistant", "content": openai_response}
             yield "", openai_messages, anthropic_messages, gemini_messages
-        if openai_response.strip() and not openai_response.startswith("Error:"):
-            openai_messages[-1] = {"role": "assistant", "content": openai_response}
 
-    # Handle Anthropic (non-streaming)
+    # Handle Anthropic (streaming)
     if "Anthropic" in providers:
-        anthropic_messages.append({"role": "user", "content": query})
-        anthropic_response = await ask_anthropic(query, history)
-        if anthropic_response.strip() and not anthropic_response.startswith("Error:"):
-            anthropic_messages.append({"role": "assistant", "content": anthropic_response})
+        async for chunk in ask_anthropic(query, history):
+            anthropic_response += chunk
+            # Update Anthropic chatbot with streaming response
+            anthropic_messages[-1] = {"role": "assistant", "content": anthropic_response}
+            yield "", openai_messages, anthropic_messages, gemini_messages
 
-    # Handle Gemini (non-streaming)
+    # Handle Gemini (streaming)
     if "Gemini" in providers:
-        gemini_messages.append({"role": "user", "content": query})
-        gemini_response = await ask_gemini(query, history)
-        if gemini_response.strip() and not gemini_response.startswith("Error:"):
-            gemini_messages.append({"role": "assistant", "content": gemini_response})
+        async for chunk in ask_gemini(query, history):
+            gemini_response += chunk
+            # Update Gemini chatbot with streaming response
+            gemini_messages[-1] = {"role": "assistant", "content": gemini_response}
+            yield "", openai_messages, anthropic_messages, gemini_messages
 
     # Combine responses for history
     responses = []