Spaces:
Runtime error
Runtime error
| import os | |
| import httpx | |
| import json | |
| 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: | |
| logger.error("OpenAI API key not provided") | |
| yield "Error: OpenAI API key not provided." | |
| return | |
| messages = [] | |
| for msg in history: | |
| messages.append({"role": "user", "content": msg["user"]}) | |
| if msg.get("openai"): | |
| messages.append({"role": "assistant", "content": msg["openai"]}) | |
| messages.append({"role": "user", "content": query}) | |
| headers = { | |
| "Authorization": f"Bearer {openai_api_key}", | |
| "Content-Type": "application/json" | |
| } | |
| payload = { | |
| "model": "gpt-3.5-turbo", | |
| "messages": messages, | |
| "stream": True | |
| } | |
| try: | |
| async with httpx.AsyncClient() as client: | |
| async with client.stream("POST", "https://api.openai.com/v1/chat/completions", headers=headers, json=payload) as response: | |
| response.raise_for_status() | |
| buffer = "" | |
| async for chunk in response.aiter_text(): | |
| if chunk: | |
| buffer += chunk | |
| while "\n" in buffer: | |
| line, buffer = buffer.split("\n", 1) | |
| if line.startswith("data: "): | |
| data = line[6:] | |
| if data.strip() == "[DONE]": | |
| break | |
| if not data.strip(): | |
| continue | |
| try: | |
| json_data = json.loads(data) | |
| delta = json_data["choices"][0].get("delta", {}) | |
| if "content" in delta: | |
| yield delta["content"] | |
| except Exception as e: | |
| logger.error(f"OpenAI parse error: {e}") | |
| yield f"[OpenAI Error]: {e}" | |
| except Exception as 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: | |
| logger.error("Anthropic API key not provided") | |
| yield "Error: Anthropic API key not provided." | |
| return | |
| messages = [] | |
| for msg in history: | |
| messages.append({"role": "user", "content": msg["user"]}) | |
| if msg.get("anthropic"): | |
| messages.append({"role": "assistant", "content": msg["anthropic"]}) | |
| messages.append({"role": "user", "content": query}) | |
| headers = { | |
| "x-api-key": anthropic_api_key, | |
| "anthropic-version": "2023-06-01", | |
| "Content-Type": "application/json" | |
| } | |
| payload = { | |
| "model": "claude-3-5-sonnet-20241022", | |
| "max_tokens": 1024, | |
| "messages": messages, | |
| "stream": True | |
| } | |
| try: | |
| 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 = "" | |
| async for chunk in response.aiter_text(): | |
| if chunk: | |
| buffer += chunk | |
| while "\n" in buffer: | |
| line, buffer = buffer.split("\n", 1) | |
| if line.startswith("data: "): | |
| data = line[6:] | |
| 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: | |
| yield json_data["delta"].get("text", "") | |
| except Exception as e: | |
| logger.error(f"Anthropic parse error: {e}") | |
| yield f"[Anthropic Error]: {e}" | |
| except Exception as 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: | |
| logger.error("Gemini API key not provided") | |
| yield "Error: Gemini API key not provided." | |
| return | |
| history_text = "" | |
| for msg in history: | |
| 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_prompt}]}] | |
| } | |
| try: | |
| 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}", | |
| headers=headers, | |
| json=payload | |
| ) as response: | |
| response.raise_for_status() | |
| buffer = "" | |
| async for chunk in response.aiter_text(): | |
| 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 API error: {e}") | |
| yield f"[Gemini Error]: {e}" | |