RAG Agent

1import os2import sys3import json4import traceback56try:7    import requests8except ImportError:9    import subprocess10    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])11    import requests1213# Environment14EMBLEMA_API_BASE_URL = os.getenv("EMBLEMA_API_BASE_URL", "http://localhost:3000")15USER_TOKEN = os.getenv("USER_TOKEN")16AGENT_EXECUTE_TIMEOUT = int(os.getenv("AGENT_EXECUTE_TIMEOUT", "300"))  # 5 minutes171819def execute_agent(agent_id, prompt, context_item_ids):20    """Call the agent execute API endpoint."""21    url = f"{EMBLEMA_API_BASE_URL}/api/v2/agents/{agent_id}/execute"2223    headers = {"Content-Type": "application/json"}24    if USER_TOKEN:25        headers["Authorization"] = f"Bearer {USER_TOKEN}"2627    payload = {28        "agentId": agent_id,29        "prompt": prompt,30        "contextItemIds": context_item_ids,31    }3233    print(f"Calling RAG agent: {agent_id}", file=sys.stderr)34    print(f"Prompt: {prompt[:200]}{'...' if len(prompt) > 200 else ''}", file=sys.stderr)35    print(f"Context items: {context_item_ids}", file=sys.stderr)3637    response = requests.post(38        url,39        headers=headers,40        json=payload,41        timeout=AGENT_EXECUTE_TIMEOUT,42    )4344    if response.status_code != 200:45        try:46            error_detail = response.json()47            error_msg = error_detail.get("message", response.text)48            error_code = error_detail.get("errorCode", "UNKNOWN")49        except Exception:50            error_msg = response.text51            error_code = "HTTP_ERROR"52        raise RuntimeError(53            f"Agent API returned {response.status_code} ({error_code}): {error_msg}"54        )5556    return response.json()575859def main():60    try:61        input_json = sys.stdin.read()62        execution_input = json.loads(input_json)63        inputs = execution_input.get("inputs", {})6465        agent_id = "rag-qwen3.5-35b"66        prompt = inputs.get("prompt", "")67        drive_item_id = inputs.get("driveItemId", "")6869        if not prompt or not prompt.strip():70            raise ValueError("Prompt is required")7172        if not drive_item_id or not drive_item_id.strip():73            raise ValueError("Drive Item is required - select a folder or file from Drive")7475        # Wrap single drive item ID into array for the API76        context_item_ids = [drive_item_id.strip()]7778        # Execute RAG agent79        result = execute_agent(agent_id, prompt, context_item_ids)8081        text = result.get("text", "")82        tool_results = result.get("toolResults", [])83        steps = result.get("steps", 0)8485        print(f"RAG agent completed in {steps} steps", file=sys.stderr)86        print(f"Text length: {len(text)} chars", file=sys.stderr)87        print(f"Tool results: {len(tool_results)} items", file=sys.stderr)8889        # Output90        output = {91            "text": text,92            "toolResults": tool_results,93        }94        print(json.dumps(output, indent=2, ensure_ascii=False))9596    except Exception as e:97        error_output = {98            "error": str(e),99            "errorType": type(e).__name__,100            "traceback": traceback.format_exc(),101        }102        print(json.dumps(error_output), file=sys.stderr)103        sys.exit(1)104105106if __name__ == "__main__":107    main()