Spaces:
Running
Running
| from agent import AgentContext, UserMessage | |
| from python.helpers.api import ApiHandler, Request, Response | |
| from python.helpers import files, dotenv | |
| from initialize import initialize_agent | |
| import os | |
| import json | |
| import base64 | |
| import queue | |
| import traceback | |
| class Stream(ApiHandler): | |
| def requires_auth(cls) -> bool: | |
| return False | |
| def requires_csrf(cls) -> bool: | |
| return False | |
| def requires_api_key(cls) -> bool: | |
| return True | |
| async def process(self, input: dict, request: Request) -> Response: | |
| try: | |
| text = input.get("message") or input.get("text") or "" | |
| ctxid = input.get("context") | |
| subagent = input.get("subagent") or input.get("profile") | |
| file_data = input.get("file") | |
| file_name = input.get("file_name", "uploaded_file") | |
| dotenv.load_dotenv() | |
| # Automatically use BLABLADOR_API_KEY for 'other' provider if available | |
| blablador_key = os.getenv("BLABLADOR_API_KEY") | |
| if blablador_key: | |
| os.environ.setdefault("OTHER_API_KEY", blablador_key) | |
| os.environ.setdefault("API_KEY_OTHER", blablador_key) | |
| context = self.get_context(ctxid) | |
| config = initialize_agent() | |
| if config.chat_model.provider == "Other OpenAI compatible": | |
| config.chat_model.provider = "other" | |
| if config.utility_model.provider == "Other OpenAI compatible": | |
| config.utility_model.provider = "other" | |
| if subagent: | |
| config.profile = subagent | |
| if subagent not in config.knowledge_subdirs: | |
| config.knowledge_subdirs.append(subagent) | |
| context.config = config | |
| curr_agent = context.agent0 | |
| while curr_agent: | |
| curr_agent.config = config | |
| curr_agent = curr_agent.data.get(curr_agent.DATA_NAME_SUBORDINATE) | |
| attachment_paths = [] | |
| if file_data: | |
| # Sanitize file name to prevent path traversal | |
| file_name = os.path.basename(file_name) | |
| knowledge_dir = files.get_abs_path("knowledge/custom") | |
| os.makedirs(knowledge_dir, exist_ok=True) | |
| save_path = os.path.join(knowledge_dir, file_name) | |
| try: | |
| if isinstance(file_data, str) and "," in file_data: | |
| header, encoded = file_data.split(",", 1) | |
| file_data = encoded | |
| decoded_data = base64.b64decode(file_data) | |
| with open(save_path, "wb") as f: | |
| f.write(decoded_data) | |
| except Exception: | |
| with open(save_path, "w") as f: | |
| f.write(str(file_data)) | |
| attachment_paths.append(save_path) | |
| sync_queue = queue.Queue() | |
| context.stream_queue = sync_queue | |
| msg = UserMessage(text, attachment_paths) | |
| task = context.communicate(msg) | |
| def generate(): | |
| try: | |
| while task.is_alive() or not sync_queue.empty(): | |
| try: | |
| chunk = sync_queue.get(timeout=0.1) | |
| yield f"data: {json.dumps(chunk)}\n\n" | |
| except queue.Empty: | |
| if not task.is_alive(): | |
| break | |
| continue | |
| except Exception as e: | |
| yield f"data: {json.dumps({'type': 'error', 'text': str(e)})}\n\n" | |
| break | |
| try: | |
| result = task.result_sync(timeout=300) | |
| yield f"data: {json.dumps({'type': 'final', 'text': result})}\n\n" | |
| except Exception as e: | |
| yield f"data: {json.dumps({'type': 'error', 'text': f'Result error: {str(e)}'})}\n\n" | |
| except Exception as e: | |
| yield f"data: {json.dumps({'type': 'error', 'text': f'Generator error: {str(e)}'})}\n\n" | |
| finally: | |
| if hasattr(context, 'stream_queue'): | |
| delattr(context, 'stream_queue') | |
| return Response(generate(), mimetype='text/event-stream') | |
| except Exception as e: | |
| return Response(f"Error: {str(e)}\n{traceback.format_exc()}", status=500) | |