chunks.json

[
    {
        "id": "build.py::2",
        "input_type": "file",
        "content": "import os\nimport subprocess\nimport shutil\nimport toml\nimport sys\n\ndef run_command(command, cwd=None):\n    result = subprocess.run(command, shell=True, cwd=cwd, check=True)\n    return result\n\ndef npm_install():\n    print(\"Running npm install\")\n    run_command(\"npm install\", cwd=\"ell-studio\")\n\n\ndef npm_build():\n    print(\"Running npm build\")\n    run_command(\"npm run build\", cwd=\"ell-studio\")\n    print(\"Copying static files\")\n    source_dir = os.path.join(\"ell-studio\", \"build\")\n    target_dir = os.path.join(\"src\", \"ell\", \"studio\", \"static\")\n    shutil.rmtree(target_dir, ignore_errors=True)\n    shutil.copytree(source_dir, target_dir)\n    print(f\"Copied static files from {source_dir} to {target_dir}\")\n\n\ndef get_ell_version():\n    pyproject_path = \"pyproject.toml\"\n    pyproject_data = toml.load(pyproject_path)\n    return pyproject_data[\"tool\"][\"poetry\"][\"version\"]\n\n\ndef run_pytest():\n    print(\"Running pytest\")\n    try:\n        run_command(\"pytest\", cwd=\"tests\")\n    except subprocess.CalledProcessError:\n        print(\"Pytest failed. Aborting build.\")\n        sys.exit(1)\n\n\ndef run_all_examples():\n    print(\"Running all examples\")\n    try:\n        run_command(\"python run_all_examples.py -w 16\", cwd=\"tests\")\n    except subprocess.CalledProcessError:\n        print(\"Some examples failed. Please review the output above.\")\n        user_input = input(\"Do you want to continue with the build? (y/n): \").lower()\n        if user_input != 'y':\n            print(\"Aborting build.\")\n            sys.exit(1)\n\n\ndef main():\n    ell_version = get_ell_version()\n    os.environ['REACT_APP_ELL_VERSION'] = ell_version\n    npm_install()\n    npm_build()\n    run_pytest()\n    run_all_examples()\n    print(\"Build completed successfully.\")\n\n\nif __name__ == \"__main__\":\n    main()",
        "filepath": "build.py",
        "metadata": {
            "file_path": "build.py",
            "file_name": "build.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 427,
            "span_ids": [
                "impl",
                "npm_install",
                "run_command",
                "get_ell_version",
                "npm_build",
                "imports",
                "run_pytest",
                "run_all_examples",
                "main"
            ],
            "start_line": 1,
            "end_line": 66,
            "community": null
        },
        "node_id": "build.py::2"
    },
    {
        "id": "0.1.0\\autostreamprevention.py::1",
        "input_type": "file",
        "content": "import openai\nimport os\n\n# Define the function to stream the response\ndef stream_openai_response(prompt):\n    try:\n        # Make the API call\n        response = openai.chat.completions.create(\n            model=\"o1-mini\",  # Specify the model\n            messages=[{\"role\": \"user\", \"content\": prompt}],\n            stream=True  # Enable streaming\n        )\n\n        # Stream the response\n        for chunk in response:\n            if chunk.choices[0].delta.get(\"content\"):\n                print(chunk.choices[0].delta.content, end=\"\", flush=True)\n\n        print()  # Print a newline at the end\n\n    except Exception as e:\n        print(f\"An error occurred: {e}\")\n\n# Example usage\nprompt = \"Tell me a short joke.\"\nstream_openai_response(prompt)",
        "filepath": "docs\\ramblings\\0.1.0\\autostreamprevention.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\autostreamprevention.py",
            "file_name": "autostreamprevention.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 169,
            "span_ids": [
                "stream_openai_response",
                "impl",
                "imports"
            ],
            "start_line": 1,
            "end_line": 26,
            "community": null
        },
        "node_id": "0.1.0\\autostreamprevention.py::1"
    },
    {
        "id": "0.1.0\\cem.py::1",
        "input_type": "file",
        "content": "import gym\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nfrom gym.vector import AsyncVectorEnv\nimport random\n\n# Set random seeds for reproducibility\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\n# Hyperparameters\nNUM_ENVIRONMENTS = 4           # Reduced for simplicity\nNUM_ITERATIONS = 50            # Number of training iterations\nTRAJECTORIES_PER_ITER = 100    # Total number of trajectories per iteration\nELITE_PERCENT = 10             # Top k% trajectories to select\nLEARNING_RATE = 1e-3\nBATCH_SIZE = 64\nMAX_STEPS = 500                # Max steps per trajectory\nENV_NAME = 'CartPole-v1'",
        "filepath": "docs\\ramblings\\0.1.0\\cem.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cem.py",
            "file_name": "cem.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 168,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 23,
            "community": null
        },
        "node_id": "0.1.0\\cem.py::1"
    },
    {
        "id": "0.1.0\\cem.py::2",
        "input_type": "file",
        "content": "       # Gym environment\n\n# Define the Policy Network\nclass PolicyNetwork(nn.Module):\n    def __init__(self, state_dim, action_dim, hidden_dim=128):\n        super(PolicyNetwork, self).__init__()\n        self.fc = nn.Sequential(\n            nn.Linear(state_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Linear(hidden_dim, action_dim)\n        )\n\n    def forward(self, state):\n        logits = self.fc(state)\n        return logits\n\n    def get_action(self, state):\n        logits = self.forward(state)\n        action_probs = torch.softmax(logits, dim=-1)\n        action = torch.multinomial(action_probs, num_samples=1)\n        return action.squeeze(-1)\n\n# Function to create multiple environments\ndef make_env(env_name, seed):\n    def _init():\n        env = gym.make(env_name)\n        return env\n    return _init",
        "filepath": "docs\\ramblings\\0.1.0\\cem.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cem.py",
            "file_name": "cem.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 192,
            "span_ids": [
                "PolicyNetwork.get_action",
                "PolicyNetwork.forward",
                "imports",
                "PolicyNetwork",
                "make_env",
                "PolicyNetwork.__init__"
            ],
            "start_line": 23,
            "end_line": 52,
            "community": null
        },
        "node_id": "0.1.0\\cem.py::2"
    },
    {
        "id": "0.1.0\\cem.py::3",
        "input_type": "file",
        "content": "def collect_trajectories(envs, policy, num_trajectories, max_steps):\n    trajectories = []\n    num_envs = envs.num_envs\n\n    # Handle the return type of reset()\n    reset_output = envs.reset()\n    if isinstance(reset_output, tuple) or isinstance(reset_output, list):\n        obs = reset_output[0]  # Extract observations\n    else:\n        obs = reset_output\n\n    done_envs = [False] * num_envs\n    steps = 0\n\n    # Initialize storage for states, actions, and rewards per environment\n    env_states = [[] for _ in range(num_envs)]\n    env_actions = [[] for _ in range(num_envs)]\n    env_rewards = [0.0 for _ in range(num_envs)]\n    total_collected = 0\n    # ... other code",
        "filepath": "docs\\ramblings\\0.1.0\\cem.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cem.py",
            "file_name": "cem.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 175,
            "span_ids": [
                "collect_trajectories"
            ],
            "start_line": 54,
            "end_line": 72,
            "community": null
        },
        "node_id": "0.1.0\\cem.py::3"
    },
    {
        "id": "0.1.0\\cem.py::4",
        "input_type": "file",
        "content": "def collect_trajectories(envs, policy, num_trajectories, max_steps):\n    # ... other code\n\n    while total_collected < num_trajectories and steps < max_steps:\n        # Convert observations to tensor efficiently\n        try:\n            # Ensure 'obs' is a NumPy array\n            if not isinstance(obs, np.ndarray):\n                print(f\"Unexpected type for observations: {type(obs)}\")\n                raise ValueError(\"Observations are not a NumPy array.\")\n\n            # Convert observations to tensor using from_numpy for efficiency\n            obs_tensor = torch.from_numpy(obs).float()\n            # Ensure the observation dimension matches expected\n            assert obs_tensor.shape[1] == 4, f\"Expected observation dimension 4, got {obs_tensor.shape[1]}\"\n        except Exception as e:\n            print(f\"Error converting observations to tensor at step {steps}: {e}\")\n            print(f\"Observations: {obs}\")\n            raise e\n\n        with torch.no_grad():\n            actions = policy.get_action(obs_tensor).cpu().numpy()\n\n        # Unpack step based on Gym version\n        try:\n            # For Gym versions >=0.26, step returns five values\n            next_obs, rewards, dones, truncs, infos = envs.step(actions)\n        except ValueError:\n            # For older Gym versions, step returns four values\n            next_obs, rewards, dones, infos = envs.step(actions)\n            truncs = [False] * len(dones)  # Assume no truncations if not provided\n\n        # Handle the reset output of step()\n        if isinstance(next_obs, tuple) or isinstance(next_obs, list):\n            next_obs = next_obs[0]  # Extract observations\n\n        # Ensure infos is a list\n        if not isinstance(infos, list):\n            infos = [{} for _ in range(num_envs)]  # Default to empty dicts\n\n        for i in range(num_envs):\n            if not done_envs[i]:\n                # Check if obs[i] has the correct shape\n                if len(obs[i]) != 4:\n                    print(f\"Unexpected observation shape for env {i}: {obs[i]}\")\n                    continue  # Skip this step for the problematic environment\n\n                env_states[i].append(obs[i])\n                env_actions[i].append(actions[i])\n                env_rewards[i] += rewards[i]\n                if dones[i] or truncs[i]:\n                    # Extract reward from infos\n                    if isinstance(infos[i], dict):\n                        episode_info = infos[i].get('episode', {})\n                        traj_reward = episode_info.get('r') if 'r' in episode_info else env_rewards[i]\n                    else:\n                        # Handle cases where infos[i] is not a dict\n                        traj_reward = env_rewards[i]\n                        print(f\"Warning: infos[{i}] is not a dict. Received type: {type(infos[i])}\")\n\n                    trajectories.append({\n                        'states': env_states[i],\n                        'actions': env_actions[i],\n                        'reward': traj_reward\n                    })\n                    total_collected += 1\n                    env_states[i] = []\n                    env_actions[i] = []\n                    env_rewards[i] = 0.0\n                    done_envs[i] = True\n\n        obs = next_obs\n        steps += 1\n\n        # Reset environments that are done\n        if any(done_envs):\n            indices = [i for i, done in enumerate(done_envs) if done]\n            if total_collected < num_trajectories:\n                for i in indices:\n                    try:\n                        # Directly reset the environment\n                        reset_output = envs.envs[i].reset()\n                        if isinstance(reset_output, tuple) or isinstance(reset_output, list):\n                            # For Gym versions where reset returns (obs, info)\n                            obs[i] = reset_output[0]\n                        else:\n                            # For Gym versions where reset returns only obs\n                            obs[i] = reset_output\n                        done_envs[i] = False\n                    except Exception as e:\n                        print(f\"Error resetting environment {i}: {e}\")\n                        # Optionally, handle the failure (e.g., retry, terminate the environment)\n                        done_envs[i] = False  # Prevent infinite loop\n\n    return trajectories",
        "filepath": "docs\\ramblings\\0.1.0\\cem.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cem.py",
            "file_name": "cem.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 872,
            "span_ids": [
                "collect_trajectories"
            ],
            "start_line": 74,
            "end_line": 165,
            "community": null
        },
        "node_id": "0.1.0\\cem.py::4"
    },
    {
        "id": "0.1.0\\cem.py::5",
        "input_type": "file",
        "content": "def select_elite(trajectories, percentile=ELITE_PERCENT):\n    rewards = [traj['reward'] for traj in trajectories]\n    if not rewards:\n        return []\n    reward_threshold = np.percentile(rewards, 100 - percentile)\n    elite_trajectories = [traj for traj in trajectories if traj['reward'] >= reward_threshold]\n    return elite_trajectories\n\n# Function to create training dataset from elite trajectories\ndef create_training_data(elite_trajectories):\n    states = []\n    actions = []\n    for traj in elite_trajectories:\n        states.extend(traj['states'])\n        actions.extend(traj['actions'])\n    if not states or not actions:\n        return None, None\n    # Convert lists to NumPy arrays first for efficiency\n    states = np.array(states, dtype=np.float32)\n    actions = np.array(actions, dtype=np.int64)\n    # Convert to PyTorch tensors\n    states = torch.from_numpy(states)\n    actions = torch.from_numpy(actions)\n    return states, actions",
        "filepath": "docs\\ramblings\\0.1.0\\cem.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cem.py",
            "file_name": "cem.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 209,
            "span_ids": [
                "create_training_data",
                "select_elite"
            ],
            "start_line": 168,
            "end_line": 191,
            "community": null
        },
        "node_id": "0.1.0\\cem.py::5"
    },
    {
        "id": "0.1.0\\cem.py::6",
        "input_type": "file",
        "content": "# Main execution code\nif __name__ == '__main__':\n    # Initialize environments\n    env_fns = [make_env(ENV_NAME, SEED + i) for i in range(NUM_ENVIRONMENTS)]\n    envs = AsyncVectorEnv(env_fns)\n\n    # Get environment details\n    dummy_env = gym.make(ENV_NAME)\n    state_dim = dummy_env.observation_space.shape[0]\n    action_dim = dummy_env.action_space.n\n    dummy_env.close()\n\n    # Initialize policy network and optimizer\n    policy = PolicyNetwork(state_dim, action_dim)\n    optimizer = optim.Adam(policy.parameters(), lr=LEARNING_RATE)\n    criterion = nn.CrossEntropyLoss()\n\n    # Training Loop\n    for iteration in range(1, NUM_ITERATIONS + 1):\n        try:\n            # Step 1: Collect Trajectories\n            trajectories = collect_trajectories(envs, policy, TRAJECTORIES_PER_ITER, MAX_STEPS)\n        except Exception as e:\n            print(f\"Error during trajectory collection at iteration {iteration}: {e}\")\n            break\n\n        # Step 2: Select Elite Trajectories\n        elite_trajectories = select_elite(trajectories, ELITE_PERCENT)\n\n        if len(elite_trajectories) == 0:\n            print(f\"Iteration {iteration}: No elite trajectories found. Skipping update.\")\n            continue\n\n        # Step 3: Create Training Data\n        states, actions = create_training_data(elite_trajectories)\n\n        if states is None or actions is None:\n            print(f\"Iteration {iteration}: No training data available. Skipping update.\")\n            continue\n\n        # Step 4: Behavioral Cloning (Policy Update)\n        dataset_size = states.size(0)\n        indices = np.arange(dataset_size)\n        np.random.shuffle(indices)\n\n        for start in range(0, dataset_size, BATCH_SIZE):\n            end = start + BATCH_SIZE\n            batch_indices = indices[start:end]\n            batch_states = states[batch_indices]\n            batch_actions = actions[batch_indices]\n\n            optimizer.zero_grad()\n            logits = policy(batch_states)\n            loss = criterion(logits, batch_actions)\n            loss.backward()\n            optimizer.step()\n\n        # Step 5: Evaluate Current Policy\n        avg_reward = np.mean([traj['reward'] for traj in elite_trajectories])\n        print(f\"Iteration {iteration}: Elite Trajectories: {len(elite_trajectories)}, Average Reward: {avg_reward:.2f}\")\n\n    # Close environments\n    envs.close()\n\n    # Testing the Trained Policy\n    def test_policy(policy, env_name=ENV_NAME, episodes=5, max_steps=500):\n        env = gym.make(env_name)\n        total_rewards = []\n        for episode in range(episodes):\n            obs, _ = env.reset()\n            done = False\n            episode_reward = 0\n            for _ in range(max_steps):\n                obs_tensor = torch.from_numpy(obs).float().unsqueeze(0)\n                with torch.no_grad():\n                    action = policy.get_action(obs_tensor).item()\n                obs, reward, done, info, _ = env.step(action)\n                episode_reward += reward\n                if done:\n                    break\n            total_rewards.append(episode_reward)\n            print(f\"Test Episode {episode + 1}: Reward: {episode_reward}\")\n        env.close()\n        print(f\"Average Test Reward over {episodes} episodes: {np.mean(total_rewards):.2f}\")\n\n    # Run the test\n    test_policy(policy)",
        "filepath": "docs\\ramblings\\0.1.0\\cem.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cem.py",
            "file_name": "cem.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 717,
            "span_ids": [
                "create_training_data",
                "impl:22"
            ],
            "start_line": 193,
            "end_line": 279,
            "community": null
        },
        "node_id": "0.1.0\\cem.py::6"
    },
    {
        "id": "0.1.0\\context_versioning.py::1",
        "input_type": "file",
        "content": "import inspect\nimport ast\nfrom contextlib import contextmanager\n\n@contextmanager\ndef context():\n    # Get the current frame\n    frame = inspect.currentframe()\n    try:\n        # Get the caller's frame\n        caller_frame = frame.f_back.f_back\n        # Get the filename and line number where the context manager is called\n        filename = caller_frame.f_code.co_filename\n        lineno = caller_frame.f_lineno\n\n        # Read the source code from the file\n        with open(filename, 'r') as f:\n            source = f.read()\n\n        # Parse the source code into an AST\n        parsed = ast.parse(source, filename)\n        # print(source)\n        # Find the 'with' statement at the given line number\n        class WithVisitor(ast.NodeVisitor):\n            def __init__(self, target_lineno):\n                self.target_lineno = target_lineno\n                self.with_node = None\n\n            def visit_With(self, node):\n                if node.lineno <= self.target_lineno <= node.end_lineno:\n                    self.with_node = node\n                self.generic_visit(node)\n\n        visitor = WithVisitor(lineno)\n        visitor.visit(parsed)\n\n        # print(parsed, source)\n        if visitor.with_node:\n            # Extract the source code of the block inside 'with'\n            start = visitor.with_node.body[0].lineno\n            end = visitor.with_node.body[-1].end_lineno\n            block_source = '\\n'.join(source.splitlines()[start-1:end])\n            print(\"Source code inside 'with' block:\")\n            print(block_source)\n        else:\n            print(\"Could not find the 'with' block.\")\n\n        # Yield control to the block inside 'with'\n        yield\n    finally:\n        # Any cleanup can be done here\n        pass\n\nfrom context_versioning import context\n# Example usage\nif __name__ == \"__main__\":\n    with context():\n        x = 10\n        y = x * 2\n        print(y)",
        "filepath": "docs\\ramblings\\0.1.0\\context_versioning.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\context_versioning.py",
            "file_name": "context_versioning.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 405,
            "span_ids": [
                "impl",
                "imports",
                "context"
            ],
            "start_line": 2,
            "end_line": 63,
            "community": null
        },
        "node_id": "0.1.0\\context_versioning.py::1"
    },
    {
        "id": "0.1.0\\cpbo.py::1",
        "input_type": "file",
        "content": "import gym\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nfrom collections import namedtuple\nfrom torch.utils.data import DataLoader, TensorDataset\n\n# Define a simple policy network\nclass PolicyNetwork(nn.Module):\n    def __init__(self, state_dim, action_dim, hidden_dim=128):\n        super(PolicyNetwork, self).__init__()\n        self.network = nn.Sequential(\n            nn.Linear(state_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Linear(hidden_dim, hidden_dim),\n            nn.ReLU(),\n            nn.Linear(hidden_dim, action_dim),\n            nn.Softmax(dim=-1)  # Output action probabilities\n        )\n\n    def forward(self, x):\n        return self.network(x)",
        "filepath": "docs\\ramblings\\0.1.0\\cpbo.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cpbo.py",
            "file_name": "cpbo.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 151,
            "span_ids": [
                "PolicyNetwork.forward",
                "PolicyNetwork.__init__",
                "imports",
                "PolicyNetwork"
            ],
            "start_line": 1,
            "end_line": 23,
            "community": null
        },
        "node_id": "0.1.0\\cpbo.py::1"
    },
    {
        "id": "0.1.0\\cpbo.py::2",
        "input_type": "file",
        "content": "# Function to collect trajectories\ndef collect_trajectories(env, policy, num_episodes, device):\n    trajectories = []\n    Episode = namedtuple('Episode', ['states', 'actions', 'rewards'])\n\n    for episode_num in range(num_episodes):\n        states = []\n        actions = []\n        rewards = []\n        # Handle Gym's updated reset() API\n        state, info = env.reset(seed=42 + episode_num)  # Optional: set seed for reproducibility\n        done = False\n\n        while not done:\n            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)\n            with torch.no_grad():\n                action_probs = policy(state_tensor)\n            action_dist = torch.distributions.Categorical(action_probs)\n            action = action_dist.sample().item()\n\n            # Handle Gym's updated step() API\n            next_state, reward, terminated, truncated, info = env.step(action)\n            done = terminated or truncated\n\n            states.append(state)\n            actions.append(action)\n            rewards.append(reward)\n\n            state = next_state\n\n        trajectories.append(Episode(states, actions, rewards))\n\n    return trajectories\n\n# Function to compute returns\ndef compute_returns(trajectories, gamma=0.99):\n    all_returns = []\n    for episode in trajectories:\n        returns = []\n        G = 0\n        for reward in reversed(episode.rewards):\n            G = reward + gamma * G\n            returns.insert(0, G)\n        all_returns.extend(returns)\n    return all_returns",
        "filepath": "docs\\ramblings\\0.1.0\\cpbo.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cpbo.py",
            "file_name": "cpbo.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 305,
            "span_ids": [
                "PolicyNetwork.forward",
                "collect_trajectories",
                "compute_returns"
            ],
            "start_line": 25,
            "end_line": 69,
            "community": null
        },
        "node_id": "0.1.0\\cpbo.py::2"
    },
    {
        "id": "0.1.0\\cpbo.py::3",
        "input_type": "file",
        "content": "# Function to create labeled dataset\ndef create_labeled_dataset(trajectories, gamma=0.99, device='cpu'):\n    states = []\n    actions = []\n    labels = []\n\n    all_returns = compute_returns(trajectories, gamma)\n    all_returns = np.array(all_returns)\n    median_return = np.median(all_returns)\n\n    for episode in trajectories:\n        for t in range(len(episode.rewards)):\n            # Compute return from timestep t\n            G = sum([gamma**k * episode.rewards[t + k] for k in range(len(episode.rewards) - t)])\n            label = 1 if G >= median_return else 0\n            states.append(episode.states[t])\n            actions.append(episode.actions[t])\n            labels.append(label)\n\n    # Convert lists to NumPy arrays first for efficiency\n    states = np.array(states)\n    actions = np.array(actions)\n    labels = np.array(labels)\n\n    # Convert to PyTorch tensors\n    states = torch.FloatTensor(states).to(device)\n    actions = torch.LongTensor(actions).to(device)\n    labels = torch.FloatTensor(labels).to(device)\n\n    return states, actions, labels",
        "filepath": "docs\\ramblings\\0.1.0\\cpbo.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cpbo.py",
            "file_name": "cpbo.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 237,
            "span_ids": [
                "create_labeled_dataset",
                "compute_returns"
            ],
            "start_line": 71,
            "end_line": 100,
            "community": null
        },
        "node_id": "0.1.0\\cpbo.py::3"
    },
    {
        "id": "0.1.0\\cpbo.py::4",
        "input_type": "file",
        "content": "# Function to perform behavioral cloning update\ndef behavioral_cloning_update(policy, optimizer, dataloader, device):\n    criterion = nn.BCELoss()\n    policy.train()\n\n    for states, actions, labels in dataloader:\n        optimizer.zero_grad()\n        action_probs = policy(states)\n        # Gather the probability of the taken action\n        selected_probs = action_probs.gather(1, actions.unsqueeze(1)).squeeze(1)\n        # Labels are 1 for good actions, 0 for bad actions\n        loss = criterion(selected_probs, labels)\n        loss.backward()\n        optimizer.step()",
        "filepath": "docs\\ramblings\\0.1.0\\cpbo.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cpbo.py",
            "file_name": "cpbo.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 122,
            "span_ids": [
                "create_labeled_dataset",
                "behavioral_cloning_update"
            ],
            "start_line": 102,
            "end_line": 115,
            "community": null
        },
        "node_id": "0.1.0\\cpbo.py::4"
    },
    {
        "id": "0.1.0\\cpbo.py::5",
        "input_type": "file",
        "content": "# Evaluation function\ndef evaluate_policy(env, policy, device, episodes=5):\n    policy.eval()\n    total_rewards = []\n    for _ in range(episodes):\n        state, info = env.reset()\n        done = False\n        ep_reward = 0\n        while not done:\n            state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)\n            with torch.no_grad():\n                action_probs = policy(state_tensor)\n            action = torch.argmax(action_probs, dim=1).item()\n            # Handle Gym's updated step() API\n            next_state, reward, terminated, truncated, info = env.step(action)\n            done = terminated or truncated\n            ep_reward += reward\n            state = next_state\n        total_rewards.append(ep_reward)\n    average_reward = np.mean(total_rewards)\n    return average_reward",
        "filepath": "docs\\ramblings\\0.1.0\\cpbo.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cpbo.py",
            "file_name": "cpbo.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 167,
            "span_ids": [
                "behavioral_cloning_update",
                "evaluate_policy"
            ],
            "start_line": 117,
            "end_line": 137,
            "community": null
        },
        "node_id": "0.1.0\\cpbo.py::5"
    },
    {
        "id": "0.1.0\\cpbo.py::6",
        "input_type": "file",
        "content": "# Main CBPO algorithm\ndef CBPO(env_name='CartPole-v1', num_epochs=10, num_episodes_per_epoch=100, gamma=0.99, \n         batch_size=64, learning_rate=1e-3, device='cpu'):\n\n    env = gym.make(env_name)\n    state_dim = env.observation_space.shape[0]\n    action_dim = env.action_space.n\n\n    policy = PolicyNetwork(state_dim, action_dim).to(device)\n    optimizer = optim.Adam(policy.parameters(), lr=learning_rate)\n\n    for epoch in range(num_epochs):\n        print(f\"Epoch {epoch+1}/{num_epochs}\")\n\n        # 1. Collect trajectories\n        trajectories = collect_trajectories(env, policy, num_episodes_per_epoch, device)\n\n        # 2. Create labeled dataset\n        states, actions, labels = create_labeled_dataset(trajectories, gamma, device)\n\n        # 3. Create DataLoader\n        dataset = TensorDataset(states, actions, labels)\n        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)\n\n        # 4. Behavioral Cloning Update\n        behavioral_cloning_update(policy, optimizer, dataloader, device)\n\n        # 5. Evaluate current policy\n        avg_reward = evaluate_policy(env, policy, device)\n        print(f\"Average Reward: {avg_reward}\")\n\n        # Early stopping if solved\n        if avg_reward >= env.spec.reward_threshold:\n            print(f\"Environment solved in {epoch+1} epochs!\")\n            break\n\n    env.close()\n    return policy",
        "filepath": "docs\\ramblings\\0.1.0\\cpbo.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cpbo.py",
            "file_name": "cpbo.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 319,
            "span_ids": [
                "CBPO",
                "evaluate_policy"
            ],
            "start_line": 139,
            "end_line": 176,
            "community": null
        },
        "node_id": "0.1.0\\cpbo.py::6"
    },
    {
        "id": "0.1.0\\cpbo.py::7",
        "input_type": "file",
        "content": "if __name__ == \"__main__\":\n    # Check if GPU is available\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    print(f\"Using device: {device}\")\n\n    # Run CBPO\n    trained_policy = CBPO(\n        env_name='CartPole-v1',\n        num_epochs=50,\n        num_episodes_per_epoch=500,\n        gamma=0.99,\n        batch_size=64,\n        learning_rate=1e-3,\n        device=device\n    )\n\n    # Final Evaluation\n    env = gym.make('CartPole-v1')\n    final_avg_reward = evaluate_policy(env, trained_policy, device, episodes=20)\n    print(f\"Final Average Reward over 20 episodes: {final_avg_reward}\")\n    env.close()\n\n    # Save the trained policy\n    torch.save(trained_policy.state_dict(), \"trained_cartpole_policy.pth\")\n    print(\"Trained policy saved to trained_cartpole_policy.pth\")\n\n    # Demo the trained policy with rendering\n    env = gym.make('CartPole-v1', render_mode='human')\n    state, _ = env.reset()\n    done = False\n    total_reward = 0\n\n    while not done:\n        state_tensor = torch.FloatTensor(state).unsqueeze(0).to(device)\n        action = trained_policy(state_tensor).argmax().item()\n        state, reward, terminated, truncated, _ = env.step(action)\n        total_reward += reward\n        done = terminated or truncated\n        env.render()\n\n    print(f\"Demo episode finished with total reward: {total_reward}\")\n    env.close()",
        "filepath": "docs\\ramblings\\0.1.0\\cpbo.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\cpbo.py",
            "file_name": "cpbo.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 333,
            "span_ids": [
                "impl"
            ],
            "start_line": 178,
            "end_line": 219,
            "community": null
        },
        "node_id": "0.1.0\\cpbo.py::7"
    },
    {
        "id": "0.1.0\\metapromptingtorch.py::1",
        "input_type": "file",
        "content": "import torch as th\n\n\nweights = th.nn.Parameter(th.randn(10))\n\n\ndef forward(x):\n    return x * weights\n\n\nx = th.randn(10)\n\nprint(forward(x))\nprint(weights)\n\n# OOOH WAHT IF WE DID MANY TYPES OF LEARNABLES in",
        "filepath": "docs\\ramblings\\0.1.0\\metapromptingtorch.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\metapromptingtorch.py",
            "file_name": "metapromptingtorch.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 57,
            "span_ids": [
                "forward",
                "impl:3",
                "imports"
            ],
            "start_line": 3,
            "end_line": 18,
            "community": null
        },
        "node_id": "0.1.0\\metapromptingtorch.py::1"
    },
    {
        "id": "0.1.0\\mypytest.py::1",
        "input_type": "file",
        "content": "from typing import TypedDict\n\n\nclass Test(TypedDict):\n    name: str\n    age: int\n\n\ndef test(**t: Test):\n    print(t)\n\n# no type hinting like ts thats unfortunate.\ntest( )",
        "filepath": "docs\\ramblings\\0.1.0\\mypytest.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\mypytest.py",
            "file_name": "mypytest.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 46,
            "span_ids": [
                "Test",
                "test",
                "impl",
                "imports"
            ],
            "start_line": 1,
            "end_line": 14,
            "community": null
        },
        "node_id": "0.1.0\\mypytest.py::1"
    },
    {
        "id": "0.1.0\\test.py::1",
        "input_type": "file",
        "content": "from typing import Callable\n\n# The follwoing works...\n\n\n\ndef decorator(fn : Callable):\n    def wrapper(*args, **kwargs):\n        print(\"before\")\n        result = fn(*args, **kwargs)\n        print(\"after\")\n        return result\n    return wrapper\n\n\nclass TestCallable:\n    def __init__(self, fn : Callable):\n        self.fn = fn\n\n    def __call__(self, *args, **kwargs):\n        return self.fn(*args, **kwargs)\n\ndef convert_to_test_callable(fn : Callable):\n    return TestCallable(fn)\n\nx = TestCallable(lambda : 1)\n\n@decorator\n@convert_to_test_callable\ndef test():\n    print(\"test\")\n\n@decorator\nclass MyCallable:\n    def __init__(self, fn : Callable):\n        self.fn = fn\n\n    def __call__(self, *args, **kwargs):\n        return self.fn(*args, **kwargs)\n\n# Oh so now ell.simples can actually be used as decorators on classes\r",
        "filepath": "docs\\ramblings\\0.1.0\\test.py",
        "metadata": {
            "file_path": "docs\\ramblings\\0.1.0\\test.py",
            "file_name": "test.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 208,
            "span_ids": [
                "impl",
                "test",
                "MyCallable",
                "TestCallable.__init__",
                "MyCallable.__call__",
                "decorator",
                "MyCallable.__init__",
                "TestCallable",
                "imports",
                "TestCallable.__call__",
                "convert_to_test_callable"
            ],
            "start_line": 2,
            "end_line": 43,
            "community": null
        },
        "node_id": "0.1.0\\test.py::1"
    },
    {
        "id": "src\\conf.py::1",
        "input_type": "file",
        "content": "# Configuration file for the Sphinx documentation builder.\n#\n# For the full list of built-in configuration values, see the documentation:\n# https://www.sphinx-doc.org/en/master/usage/configuration.html\n\n# -- Project information -----------------------------------------------------\n# https://www.sphinx-doc.org/en/master/usage/configuration.html#project-information\n\nproject = 'ell'\ncopyright = '2024, William Guss'\nauthor = 'William Guss'\n\n# -- General configuration ---------------------------------------------------\n# https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration\nextensions = ['sphinx.ext.autodoc', 'sphinx.ext.napoleon', 'sphinxawesome_theme', 'sphinxcontrib.autodoc_pydantic']\n\ntemplates_path = ['_templates']\nexclude_patterns = ['_build', 'Thumbs.db', '.DS_Store']\n\nhtml_theme = \"sphinxawesome_theme\"\n\n\n# Favicon configuration\nhtml_favicon = '_static/favicon.ico'\n\n# Configure syntax highlighting for Awesome Sphinx Theme\npygments_style = \"default\"\npygments_style_dark = \"dracula\"\n\n# Additional theme configuration\r",
        "filepath": "docs\\src\\conf.py",
        "metadata": {
            "file_path": "docs\\src\\conf.py",
            "file_name": "conf.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 223,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 30,
            "community": null
        },
        "node_id": "src\\conf.py::1"
    },
    {
        "id": "src\\conf.py::2",
        "input_type": "file",
        "content": "html_theme_options = {\n    \"show_prev_next\": True,\n    \"show_scrolltop\": True,\n    \"main_nav_links\": {\n        \"Docs\": \"index\",\n        \"API Reference\": \"reference/index\",\n        \"AI Jobs Board\": \"https://jobs.ell.so\",\n    },\n    \"extra_header_link_icons\": {\n        \"Discord\": {\n        \"link\": \"https://discord.gg/vWntgU52Xb\",\n            \"icon\": \"\"\"<svg xmlns=\"http://www.w3.org/2000/svg\" viewBox=\"0 0 640 512\" height=\"18\" fill=\"currentColor\"><!--!Font Awesome Free 6.6.0 by @fontawesome - https://fontawesome.com License - https://fontawesome.com/license/free Copyright 2024 Fonticons, Inc.--><path d=\"M524.5 69.8a1.5 1.5 0 0 0 -.8-.7A485.1 485.1 0 0 0 404.1 32a1.8 1.8 0 0 0 -1.9 .9 337.5 337.5 0 0 0 -14.9 30.6 447.8 447.8 0 0 0 -134.4 0 309.5 309.5 0 0 0 -15.1-30.6 1.9 1.9 0 0 0 -1.9-.9A483.7 483.7 0 0 0 116.1 69.1a1.7 1.7 0 0 0 -.8 .7C39.1 183.7 18.2 294.7 28.4 404.4a2 2 0 0 0 .8 1.4A487.7 487.7 0 0 0 176 479.9a1.9 1.9 0 0 0 2.1-.7A348.2 348.2 0 0 0 208.1 430.4a1.9 1.9 0 0 0 -1-2.6 321.2 321.2 0 0 1 -45.9-21.9 1.9 1.9 0 0 1 -.2-3.1c3.1-2.3 6.2-4.7 9.1-7.1a1.8 1.8 0 0 1 1.9-.3c96.2 43.9 200.4 43.9 295.5 0a1.8 1.8 0 0 1 1.9 .2c2.9 2.4 6 4.9 9.1 7.2a1.9 1.9 0 0 1 -.2 3.1 301.4 301.4 0 0 1 -45.9 21.8 1.9 1.9 0 0 0 -1 2.6 391.1 391.1 0 0 0 30 48.8 1.9 1.9 0 0 0 2.1 .7A486 486 0 0 0 610.7 405.7a1.9 1.9 0 0 0 .8-1.4C623.7 277.6 590.9 167.5 524.5 69.8zM222.5 337.6c-29 0-52.8-26.6-52.8-59.2S193.1 219.1 222.5 219.1c29.7 0 53.3 26.8 52.8 59.2C275.3 311 251.9 337.6 222.5 337.6zm195.4 0c-29 0-52.8-26.6-52.8-59.2S388.4 219.1 417.9 219.1c29.7 0 53.3 26.8 52.8 59.2C470.7 311 447.5 337.6 417.9 337.6z\"/></svg>\"\"\",\n            \"type\": \"font-awesome\",\n            \"name\": \"Discord\",\n        },\n    },\n\n    \"logo_light\": \"_static/ell-wide-light.png\",\n    \"logo_dark\": \"_static/ell-wide-dark.png\",\n    \n}\n\nhtml_static_path = ['_static']\n\n\n\ntemplates_path = ['_templates']",
        "filepath": "docs\\src\\conf.py",
        "metadata": {
            "file_path": "docs\\src\\conf.py",
            "file_name": "conf.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 1044,
            "span_ids": [
                "impl:23",
                "docstring"
            ],
            "start_line": 31,
            "end_line": 58,
            "community": null
        },
        "node_id": "src\\conf.py::2"
    },
    {
        "id": "ell\\__init__.py::1",
        "input_type": "file",
        "content": "\"\"\"\nell is a Python library for language model programming (LMP). It provides a simple\nand intuitive interface for working with large language models.\n\"\"\"\n\n\nfrom ell.lmp.simple import simple\nfrom ell.lmp.tool import tool\nfrom ell.lmp.complex import complex\nfrom ell.types.message import system, user, assistant, Message, ContentBlock\nfrom ell.__version__ import __version__\n\n# Import all models\nimport ell.providers\nimport ell.models\n\n\n# Import everything from configurator\nfrom ell.configurator import *",
        "filepath": "src\\ell\\__init__.py",
        "metadata": {
            "file_path": "src\\ell\\__init__.py",
            "file_name": "__init__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 106,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 20,
            "community": null
        },
        "node_id": "ell\\__init__.py::1"
    },
    {
        "id": "ell\\__version__.py::1",
        "input_type": "file",
        "content": "try:\n    from importlib.metadata import version\nexcept ImportError:\n    from importlib_metadata import version\n\n__version__ = version(\"ell-ai\")",
        "filepath": "src\\ell\\__version__.py",
        "metadata": {
            "file_path": "src\\ell\\__version__.py",
            "file_name": "__version__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 31,
            "span_ids": [
                "impl"
            ],
            "start_line": 1,
            "end_line": 7,
            "community": null
        },
        "node_id": "ell\\__version__.py::1"
    },
    {
        "id": "ell\\configurator.py::1",
        "input_type": "file",
        "content": "from functools import lru_cache, wraps\nfrom typing import Dict, Any, Optional, Tuple, Union, Type\nimport openai\nimport logging\nfrom contextlib import contextmanager\nimport threading\nfrom pydantic import BaseModel, ConfigDict, Field\nfrom ell.store import Store\nfrom ell.provider import Provider\nfrom dataclasses import dataclass, field\n\n_config_logger = logging.getLogger(__name__)\n\n@dataclass(frozen=True)\nclass _Model:\n    name: str\n    default_client: Optional[Union[openai.Client, Any]] = None\n    #XXX: Deprecation in 0.1.0\n    #XXX: We will depreciate this when streaming is implemented. \n    # Currently we stream by default for the verbose renderer,\n    # but in the future we will not support streaming by default \n    # and stream=True must be passed which will then make API providers the\n    # single source of truth for whether or not a model supports an api parameter.\n    # This makes our implementation extremely light, only requiring us to provide\n    # a list of model names in registration.\n    supports_streaming : Optional[bool] = field(default=None)",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 243,
            "span_ids": [
                "imports",
                "_Model"
            ],
            "start_line": 1,
            "end_line": 26,
            "community": null
        },
        "node_id": "ell\\configurator.py::1"
    },
    {
        "id": "ell\\configurator.py::2",
        "input_type": "file",
        "content": "class Config(BaseModel):\n    model_config = ConfigDict(arbitrary_types_allowed=True)\n    registry: Dict[str, _Model] = Field(default_factory=dict, description=\"A dictionary mapping model names to their configurations.\")\n    verbose: bool = Field(default=False, description=\"If True, enables verbose logging.\")\n    wrapped_logging: bool = Field(default=True, description=\"If True, enables wrapped logging for better readability.\")\n    override_wrapped_logging_width: Optional[int] = Field(default=None, description=\"If set, overrides the default width for wrapped logging.\")\n    store: Optional[Store] = Field(default=None, description=\"An optional Store instance for persistence.\")\n    autocommit: bool = Field(default=False, description=\"If True, enables automatic committing of changes to the store.\")\n    lazy_versioning: bool = Field(default=True, description=\"If True, enables lazy versioning for improved performance.\")\n    default_api_params: Dict[str, Any] = Field(default_factory=dict, description=\"Default parameters for language models.\")\n    default_client: Optional[openai.Client] = Field(default=None, description=\"The default OpenAI client used when a specific model client is not found.\")\n    autocommit_model: str = Field(default=\"gpt-4o-mini\", description=\"When set, changes the default autocommit model from GPT 4o mini.\")\n    providers: Dict[Type, Provider] = Field(default_factory=dict, description=\"A dictionary mapping client types to provider classes.\")\n    def __init__(self, **data):\n        super().__init__(**data)\n        self._lock = threading.Lock()\n        self._local = threading.local()",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 335,
            "span_ids": [
                "Config.__init__",
                "Config"
            ],
            "start_line": 30,
            "end_line": 46,
            "community": null
        },
        "node_id": "ell\\configurator.py::2"
    },
    {
        "id": "ell\\configurator.py::3",
        "input_type": "file",
        "content": "class Config(BaseModel):\n\n\n    def register_model(\n        self, \n        name: str,\n        default_client: Optional[Union[openai.Client, Any]] = None,\n        supports_streaming: Optional[bool] = None\n    ) -> None:\n        \"\"\"\n        Register a model with its configuration.\n        \"\"\"\n        with self._lock:\n            # XXX: Will be deprecated in 0.1.0\n            self.registry[name] = _Model(\n                name=name,\n                default_client=default_client,\n                supports_streaming=supports_streaming\n            )",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 118,
            "span_ids": [
                "Config.register_model"
            ],
            "start_line": 49,
            "end_line": 64,
            "community": null
        },
        "node_id": "ell\\configurator.py::3"
    },
    {
        "id": "ell\\configurator.py::4",
        "input_type": "file",
        "content": "class Config(BaseModel):\n\n\n\n    @contextmanager\n    def model_registry_override(self, overrides: Dict[str, _Model]):\n        \"\"\"\n        Temporarily override the model registry with new model configurations.\n\n        :param overrides: A dictionary of model names to ModelConfig instances to override.\n        :type overrides: Dict[str, ModelConfig]\n        \"\"\"\n        if not hasattr(self._local, 'stack'):\n            self._local.stack = []\n\n        with self._lock:\n            current_registry = self._local.stack[-1] if self._local.stack else self.registry\n            new_registry = current_registry.copy()\n            new_registry.update(overrides)\n\n        self._local.stack.append(new_registry)\n        try:\n            yield\n        finally:\n            self._local.stack.pop()",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 154,
            "span_ids": [
                "Config.model_registry_override"
            ],
            "start_line": 68,
            "end_line": 88,
            "community": null
        },
        "node_id": "ell\\configurator.py::4"
    },
    {
        "id": "ell\\configurator.py::5",
        "input_type": "file",
        "content": "class Config(BaseModel):\n\n    def get_client_for(self, model_name: str) -> Tuple[Optional[openai.Client], bool]:\n        \"\"\"\n        Get the OpenAI client for a specific model name.\n\n        :param model_name: The name of the model to get the client for.\n        :type model_name: str\n        :return: The OpenAI client for the specified model, or None if not found, and a fallback flag.\n        :rtype: Tuple[Optional[openai.Client], bool]\n        \"\"\"\n        current_registry = self._local.stack[-1] if hasattr(self._local, 'stack') and self._local.stack else self.registry\n        model_config = current_registry.get(model_name)\n        fallback = False\n        if not model_config:\n            warning_message = f\"Warning: A default provider for model '{model_name}' could not be found. Falling back to default OpenAI client from environment variables.\"\n            if self.verbose:\n                from colorama import Fore, Style\n                _config_logger.warning(f\"{Fore.LIGHTYELLOW_EX}{warning_message}{Style.RESET_ALL}\")\n            else:\n                _config_logger.debug(warning_message)\n            client = self.default_client\n            fallback = True\n        else:\n            client = model_config.default_client\n        return client, fallback",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 264,
            "span_ids": [
                "Config.get_client_for"
            ],
            "start_line": 90,
            "end_line": 113,
            "community": null
        },
        "node_id": "ell\\configurator.py::5"
    },
    {
        "id": "ell\\configurator.py::6",
        "input_type": "file",
        "content": "class Config(BaseModel):\n\n    def register_provider(self, provider: Provider, client_type: Type[Any]) -> None:\n        \"\"\"\n        Register a provider class for a specific client type.\n\n        :param provider_class: The provider class to register.\n        :type provider_class: Type[Provider]\n        \"\"\"\n        assert isinstance(client_type, type), \"client_type must be a type (e.g. openai.Client), not an an instance (myclient := openai.Client()))\"\n        with self._lock:\n            self.providers[client_type] = provider",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 112,
            "span_ids": [
                "Config.register_provider"
            ],
            "start_line": 115,
            "end_line": 124,
            "community": null
        },
        "node_id": "ell\\configurator.py::6"
    },
    {
        "id": "ell\\configurator.py::7",
        "input_type": "file",
        "content": "class Config(BaseModel):\n\n    def get_provider_for(self, client: Union[Type[Any], Any]) -> Optional[Provider]:\n        \"\"\"\n        Get the provider instance for a specific client instance.\n\n        :param client: The client instance to get the provider for.\n        :type client: Any\n        :return: The provider instance for the specified client, or None if not found.\n        :rtype: Optional[Provider]\n        \"\"\"\n\n        client_type = type(client) if not isinstance(client, type) else client\n        for provider_type, provider in self.providers.items():\n            if issubclass(client_type, provider_type) or client_type == provider_type:\n                return provider\n        return None\n\n# Single* instance\n# XXX: Make a singleton\nconfig = Config()",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 157,
            "span_ids": [
                "Config.get_provider_for",
                "impl:3"
            ],
            "start_line": 126,
            "end_line": 144,
            "community": null
        },
        "node_id": "ell\\configurator.py::7"
    },
    {
        "id": "ell\\configurator.py::8",
        "input_type": "file",
        "content": "def init(\n    store: Optional[Union[Store, str]] = None,\n    verbose: bool = False,\n    autocommit: bool = True,\n    lazy_versioning: bool = True,\n    default_api_params: Optional[Dict[str, Any]] = None,\n    default_client: Optional[Any] = None,\n    autocommit_model: str = \"gpt-4o-mini\"\n) -> None:\n    \"\"\"\n    Initialize the ELL configuration with various settings.\n\n    :param verbose: Set verbosity of ELL operations.\n    :type verbose: bool\n    :param store: Set the store for ELL. Can be a Store instance or a string path for SQLiteStore.\n    :type store: Union[Store, str], optional\n    :param autocommit: Set autocommit for the store operations.\n    :type autocommit: bool\n    :param lazy_versioning: Enable or disable lazy versioning.\n    :type lazy_versioning: bool\n    :param default_api_params: Set default parameters for language models.\n    :type default_api_params: Dict[str, Any], optional\n    :param default_openai_client: Set the default OpenAI client.\n    :type default_openai_client: openai.Client, optional\n    :param autocommit_model: Set the model used for autocommitting.\n    :type autocommit_model: str\n    \"\"\"\n    # XXX: prevent double init\n    config.verbose = verbose\n    config.lazy_versioning = lazy_versioning\n\n    if isinstance(store, str):\n        from ell.stores.sql import SQLiteStore\n        config.store = SQLiteStore(store)\n    else:\n        config.store = store\n    config.autocommit = autocommit or config.autocommit\n\n    if default_api_params is not None:\n        config.default_api_params.update(default_api_params)\n\n    if default_client is not None:\n        config.default_client = default_client\n\n    if autocommit_model is not None:\n        config.autocommit_model = autocommit_model",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 420,
            "span_ids": [
                "init"
            ],
            "start_line": 146,
            "end_line": 191,
            "community": null
        },
        "node_id": "ell\\configurator.py::8"
    },
    {
        "id": "ell\\configurator.py::9",
        "input_type": "file",
        "content": "# Existing helper functions\ndef get_store() -> Union[Store, None]:\n    return config.store\n\n# Will be deprecated at 0.1.0 \n\n# You can add more helper functions here if needed\ndef register_provider(provider: Provider, client_type: Type[Any]) -> None:\n    return config.register_provider(provider, client_type)\n\n# Deprecated now (remove at 0.1.0)\ndef set_store(*args, **kwargs) -> None:\n    raise DeprecationWarning(\"The set_store function is deprecated and will be removed in a future version. Use ell.init(store=...) instead.\")",
        "filepath": "src\\ell\\configurator.py",
        "metadata": {
            "file_path": "src\\ell\\configurator.py",
            "file_name": "configurator.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 126,
            "span_ids": [
                "get_store",
                "set_store",
                "init",
                "register_provider"
            ],
            "start_line": 193,
            "end_line": 205,
            "community": null
        },
        "node_id": "ell\\configurator.py::9"
    },
    {
        "id": "lmp\\__init__.py::1",
        "input_type": "file",
        "content": "from ell.lmp.simple import simple\nfrom ell.lmp.complex import complex",
        "filepath": "src\\ell\\lmp\\__init__.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\__init__.py",
            "file_name": "__init__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 16,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 2,
            "community": null
        },
        "node_id": "lmp\\__init__.py::1"
    },
    {
        "id": "lmp\\_track.py::1",
        "input_type": "file",
        "content": "import json\nimport logging\nimport threading\nfrom ell.types import SerializedLMP, Invocation, InvocationTrace, InvocationContents\nfrom ell.types.studio import LMPType, utc_now\nfrom ell.util._warnings import _autocommit_warning\nimport ell.util.closure\nfrom ell.configurator import config\nfrom ell.types._lstr import _lstr\n\nimport inspect\n\nimport secrets\nimport time\nfrom datetime import datetime\nfrom functools import wraps\nfrom typing import Any, Callable, Dict, Iterable, Optional, OrderedDict, Tuple\n\nfrom ell.util.serialization import get_immutable_vars\nfrom ell.util.serialization import compute_state_cache_key\nfrom ell.util.serialization import prepare_invocation_params\n\nlogger = logging.getLogger(__name__)\n\n# Thread-local storage for the invocation stack\n_invocation_stack = threading.local()\n\ndef get_current_invocation() -> Optional[str]:\n    if not hasattr(_invocation_stack, 'stack'):\n        _invocation_stack.stack = []\n    return _invocation_stack.stack[-1] if _invocation_stack.stack else None\n\ndef push_invocation(invocation_id: str):\n    if not hasattr(_invocation_stack, 'stack'):\n        _invocation_stack.stack = []\n    _invocation_stack.stack.append(invocation_id)\n\ndef pop_invocation():\n    if hasattr(_invocation_stack, 'stack') and _invocation_stack.stack:\n        _invocation_stack.stack.pop()",
        "filepath": "src\\ell\\lmp\\_track.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\_track.py",
            "file_name": "_track.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 284,
            "span_ids": [
                "push_invocation",
                "get_current_invocation",
                "imports",
                "pop_invocation"
            ],
            "start_line": 1,
            "end_line": 40,
            "community": null
        },
        "node_id": "lmp\\_track.py::1"
    },
    {
        "id": "lmp\\_track.py::2",
        "input_type": "file",
        "content": "def _track(func_to_track: Callable, *, forced_dependencies: Optional[Dict[str, Any]] = None) -> Callable:\n\n    lmp_type = getattr(func_to_track, \"__ell_type__\", LMPType.OTHER)\n\n\n    # see if it exists\n    if not hasattr(func_to_track, \"_has_serialized_lmp\"):\n        func_to_track._has_serialized_lmp = False\n\n    if not hasattr(func_to_track, \"__ell_hash__\") and not config.lazy_versioning:\n        ell.util.closure.lexically_closured_source(func_to_track, forced_dependencies)\n\n\n    @wraps(func_to_track)\n    def tracked_func(*fn_args, _get_invocation_id=False, **fn_kwargs) -> str:\n        # XXX: Cache keys and global variable binding is not thread safe.\n        # Compute the invocation id and hash the inputs for serialization.\n        invocation_id = \"invocation-\" + secrets.token_hex(16)\n\n        state_cache_key : str = None\n        if not config.store:\n            return func_to_track(*fn_args, **fn_kwargs, _invocation_origin=invocation_id)[0]\n\n        parent_invocation_id = get_current_invocation()\n        try:\n            push_invocation(invocation_id)\n\n            # Convert all positional arguments to named keyword arguments\n            sig = inspect.signature(func_to_track)\n            # Filter out kwargs that are not in the function signature\n            filtered_kwargs = {k: v for k, v in fn_kwargs.items() if k in sig.parameters}\n\n            bound_args = sig.bind(*fn_args, **filtered_kwargs)\n            bound_args.apply_defaults()\n            all_kwargs = dict(bound_args.arguments)\n\n            # Get the list of consumed lmps and clean the invocation params for serialization.\n            cleaned_invocation_params, ipstr, consumes = prepare_invocation_params( all_kwargs)\n\n            try_use_cache = hasattr(func_to_track.__wrapper__, \"__ell_use_cache__\")\n\n            if  try_use_cache:\n                # Todo: add nice logging if verbose for when using a cahced invocaiton. IN a different color with thar args..\n                if not hasattr(func_to_track, \"__ell_hash__\")  and config.lazy_versioning:\n                    fn_closure, _ = ell.util.closure.lexically_closured_source(func_to_track)\n\n                # compute the state cachekey\n                state_cache_key = compute_state_cache_key(ipstr, func_to_track.__ell_closure__)\n\n                cache_store = func_to_track.__wrapper__.__ell_use_cache__\n                cached_invocations = cache_store.get_cached_invocations(func_to_track.__ell_hash__, state_cache_key)\n\n\n                if len(cached_invocations) > 0:\n                    # XXX: Fix caching.\n                    results =  [d.deserialize() for  d in cached_invocations[0].results]\n\n                    logger.info(f\"Using cached result for {func_to_track.__qualname__} with state cache key: {state_cache_key}\")\n                    if len(results) == 1:\n                        return results[0]\n                    else:\n                        return results\n                    # Todo: Unfiy this with the non-cached case. We should go through the same code pathway.\n                else:\n                    logger.info(f\"Attempted to use cache on {func_to_track.__qualname__} but it was not cached, or did not exist in the store. Refreshing cache...\")\n\n\n            _start_time = utc_now()\n\n            # XXX: thread saftey note, if I prevent yielding right here and get the global context I should be fine re: cache key problem\n\n            # get the prompt\n            (result, invocation_api_params, metadata) = (\n                (func_to_track(*fn_args, **fn_kwargs), {}, {})\n                if lmp_type == LMPType.OTHER\n                else func_to_track(*fn_args, _invocation_origin=invocation_id, **fn_kwargs, )\n                )\n            latency_ms = (utc_now() - _start_time).total_seconds() * 1000\n            usage = metadata.get(\"usage\", {\"prompt_tokens\": 0, \"completion_tokens\": 0})\n            prompt_tokens= usage.get(\"prompt_tokens\", 0) if usage else 0\n            completion_tokens= usage.get(\"completion_tokens\", 0) if usage else 0\n\n\n            #XXX: cattrs add invocation origin here recursively on all pirmitive types within a message.\n            #XXX: This will allow all objects to be traced automatically irrespective origin rather than relying on the API to do it, it will of vourse be expensive but unify track.\n            #XXX: No other code will need to consider tracking after this point.\n\n            if not hasattr(func_to_track, \"__ell_hash__\") and config.lazy_versioning:\n                ell.util.closure.lexically_closured_source(func_to_track, forced_dependencies)\n            _serialize_lmp(func_to_track)\n\n            if not state_cache_key:\n                state_cache_key = compute_state_cache_key(ipstr, func_to_track.__ell_closure__)\n\n            _write_invocation(func_to_track, invocation_id, latency_ms, prompt_tokens, completion_tokens, \n                            state_cache_key, invocation_api_params, cleaned_invocation_params, consumes, result, parent_invocation_id)\n\n            if _get_invocation_id:\n                return result, invocation_id\n            else:\n                return result\n        finally:\n            pop_invocation()\n\n\n    func_to_track.__wrapper__  = tracked_func\n    if hasattr(func_to_track, \"__ell_api_params__\"):\n        tracked_func.__ell_api_params__ = func_to_track.__ell_api_params__\n    if hasattr(func_to_track, \"__ell_params_model__\"):\n        tracked_func.__ell_params_model__ = func_to_track.__ell_params_model__\n    tracked_func.__ell_func__ = func_to_track\n    tracked_func.__ell_track = True\n\n    return tracked_func",
        "filepath": "src\\ell\\lmp\\_track.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\_track.py",
            "file_name": "_track.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 1199,
            "span_ids": [
                "_track"
            ],
            "start_line": 43,
            "end_line": 156,
            "community": null
        },
        "node_id": "lmp\\_track.py::2"
    },
    {
        "id": "lmp\\_track.py::3",
        "input_type": "file",
        "content": "def _serialize_lmp(func):\n    # Serialize deptjh first all fo the used lmps.\n    for f in func.__ell_uses__:\n        _serialize_lmp(f)\n\n    if getattr(func, \"_has_serialized_lmp\", False):\n        return\n    func._has_serialized_lmp = False\n    fn_closure = func.__ell_closure__\n    lmp_type = func.__ell_type__\n    name = func.__qualname__\n    api_params = getattr(func, \"__ell_api_params__\", None)\n\n    lmps = config.store.get_versions_by_fqn(fqn=name)\n    version = 0\n    already_in_store = any(lmp.lmp_id == func.__ell_hash__ for lmp in lmps)\n\n    if not already_in_store:\n        commit = None\n        if lmps:\n            latest_lmp = max(lmps, key=lambda x: x.created_at)\n            version = latest_lmp.version_number + 1\n            if config.autocommit:\n                # XXX: Move this out to autocommit itself.\n                if not _autocommit_warning():\n                    from ell.util.differ import write_commit_message_for_diff\n                    commit = str(write_commit_message_for_diff(\n                    f\"{latest_lmp.dependencies}\\n\\n{latest_lmp.source}\", \n                        f\"{fn_closure[1]}\\n\\n{fn_closure[0]}\")[0])\n\n        serialized_lmp = SerializedLMP(\n            lmp_id=func.__ell_hash__,\n            name=name,\n            created_at=utc_now(),\n            source=fn_closure[0],\n            dependencies=fn_closure[1],\n            commit_message=commit,\n            initial_global_vars=get_immutable_vars(fn_closure[2]),\n            initial_free_vars=get_immutable_vars(fn_closure[3]),\n            lmp_type=lmp_type,\n            api_params=api_params if api_params else None,\n            version_number=version,\n        )\n        config.store.write_lmp(serialized_lmp, [f.__ell_hash__ for f in func.__ell_uses__])\n    func._has_serialized_lmp = True",
        "filepath": "src\\ell\\lmp\\_track.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\_track.py",
            "file_name": "_track.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 434,
            "span_ids": [
                "_serialize_lmp"
            ],
            "start_line": 158,
            "end_line": 202,
            "community": null
        },
        "node_id": "lmp\\_track.py::3"
    },
    {
        "id": "lmp\\_track.py::4",
        "input_type": "file",
        "content": "def _write_invocation(func, invocation_id, latency_ms, prompt_tokens, completion_tokens, \n                     state_cache_key, invocation_api_params, cleaned_invocation_params, consumes, result, parent_invocation_id):\n\n    invocation_contents = InvocationContents(\n        invocation_id=invocation_id,\n        params=cleaned_invocation_params,\n        results=result,\n        invocation_api_params=invocation_api_params,\n        global_vars=get_immutable_vars(func.__ell_closure__[2]),\n        free_vars=get_immutable_vars(func.__ell_closure__[3])\n    )\n\n    if invocation_contents.should_externalize and config.store.has_blob_storage:\n        invocation_contents.is_external = True\n\n        # Write to the blob store \n        blob_id = config.store.blob_store.store_blob(\n            json.dumps(invocation_contents.model_dump(\n            ), default=str, ensure_ascii=False).encode('utf-8'),\n            invocation_id\n        )\n        invocation_contents = InvocationContents(\n            invocation_id=invocation_id,\n            is_external=True,\n        )\n\n    invocation = Invocation(\n        id=invocation_id,\n        lmp_id=func.__ell_hash__,\n        created_at=utc_now(),\n        latency_ms=latency_ms,\n        prompt_tokens=prompt_tokens,\n        completion_tokens=completion_tokens,\n        state_cache_key=state_cache_key,\n        used_by_id=parent_invocation_id,\n        contents=invocation_contents\n    )\n\n    config.store.write_invocation(invocation, consumes)",
        "filepath": "src\\ell\\lmp\\_track.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\_track.py",
            "file_name": "_track.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 296,
            "span_ids": [
                "_write_invocation"
            ],
            "start_line": 204,
            "end_line": 244,
            "community": null
        },
        "node_id": "lmp\\_track.py::4"
    },
    {
        "id": "lmp\\complex.py::1",
        "input_type": "file",
        "content": "from ell.configurator import config\nfrom ell.lmp._track import _track\nfrom ell.provider import EllCallParams\nfrom ell.types._lstr import _lstr\nfrom ell.types import Message, ContentBlock\nfrom ell.types.message import LMP, InvocableLM, LMPParams, MessageOrDict, _lstr_generic\nfrom ell.types.studio import LMPType\nfrom ell.util._warnings import _no_api_key_warning, _warnings\nfrom ell.util.verbosity import compute_color, model_usage_logger_pre\n\nfrom ell.util.verbosity import model_usage_logger_post_end, model_usage_logger_post_intermediate, model_usage_logger_post_start\n\nfrom functools import wraps\nfrom typing import Any, Dict, Optional, List, Callable, Tuple, Union",
        "filepath": "src\\ell\\lmp\\complex.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\complex.py",
            "file_name": "complex.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 156,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 14,
            "community": null
        },
        "node_id": "lmp\\complex.py::1"
    },
    {
        "id": "lmp\\complex.py::2",
        "input_type": "file",
        "content": "def complex(model: str, client: Optional[Any] = None, tools: Optional[List[Callable]] = None, exempt_from_tracking=False, post_callback: Optional[Callable] = None, **api_params):\n    default_client_from_decorator = client\n    default_model_from_decorator = model\n    default_api_params_from_decorator = api_params\n    def parameterized_lm_decorator(\n        prompt: LMP,\n    ) -> Callable[..., Union[List[Message], Message]]:\n        _warnings(model, prompt, default_client_from_decorator)\n\n        @wraps(prompt)\n        def model_call(\n            *prompt_args,\n            _invocation_origin : Optional[str] = None,\n            client: Optional[Any] = None,\n            api_params: Optional[Dict[str, Any]] = None,\n            lm_params: Optional[DeprecationWarning] = None,\n            **prompt_kwargs,\n        ) -> Tuple[Any, Any, Any]:\n            # XXX: Deprecation in 0.1.0\n            if lm_params:\n                raise DeprecationWarning(\"lm_params is deprecated. Use api_params instead.\")\n\n            # promt -> str\n            res = prompt(*prompt_args, **prompt_kwargs)\n            # Convert prompt into ell messages\n            messages = _get_messages(res, prompt)\n\n            # XXX: move should log to a logger.\n            should_log = not exempt_from_tracking and config.verbose\n            # Cute verbose logging.\n            if should_log: model_usage_logger_pre(prompt, prompt_args, prompt_kwargs, \"[]\", messages) #type: ignore\n\n            # Call the model.\n            # Merge API params\n            merged_api_params = {**config.default_api_params, **default_api_params_from_decorator, **(api_params or {})}\n            n = merged_api_params.get(\"n\", 1)\n            # Merge client overrides & client registry\n            merged_client = _client_for_model(model, client or default_client_from_decorator)\n            ell_call = EllCallParams(\n                # XXX: Could change behaviour of overriding ell params for dyanmic tool calls.\n                model=merged_api_params.pop(\"model\", default_model_from_decorator),\n                messages=messages,\n                client = merged_client,\n                api_params=merged_api_params,\n                tools=tools or [],\n            )\n            # Get the provider for the model\n            provider = config.get_provider_for(ell_call.client)\n            assert provider is not None, f\"No provider found for client {ell_call.client}.\"\n\n            if should_log: model_usage_logger_post_start(n)\n            with model_usage_logger_post_intermediate(n) as _logger:\n                (result, final_api_params, metadata) = provider.call(ell_call, origin_id=_invocation_origin, logger=_logger if should_log else None)\n                if isinstance(result, list) and len(result) == 1:\n                    result = result[0]\n\n            result = post_callback(result) if post_callback else result\n            if should_log:\n                model_usage_logger_post_end()\n            #\n            #  These get sent to track. This is wack.           \n            return result, final_api_params, metadata\n        # ... other code\n    # ... other code",
        "filepath": "src\\ell\\lmp\\complex.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\complex.py",
            "file_name": "complex.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 645,
            "span_ids": [
                "complex"
            ],
            "start_line": 16,
            "end_line": 77,
            "community": null
        },
        "node_id": "lmp\\complex.py::2"
    },
    {
        "id": "lmp\\complex.py::3",
        "input_type": "file",
        "content": "def complex(model: str, client: Optional[Any] = None, tools: Optional[List[Callable]] = None, exempt_from_tracking=False, post_callback: Optional[Callable] = None, **api_params):\n    def parameterized_lm_decorator(\n        prompt: LMP,\n    ) -> Callable[..., Union[List[Message], Message]]:\n        # ... other code\n\n\n\n        model_call.__ell_api_params__ = default_api_params_from_decorator #type: ignore\n        model_call.__ell_func__ = prompt #type: ignore\n        model_call.__ell_type__ = LMPType.LM #type: ignore\n        model_call.__ell_exempt_from_tracking = exempt_from_tracking #type: ignore\n\n\n        if exempt_from_tracking:\n            return model_call\n        else:\n            # XXX: Analyze decorators with AST instead.\n            return _track(model_call, forced_dependencies=dict(tools=tools, response_format=api_params.get(\"response_format\", {})))\n    return parameterized_lm_decorator",
        "filepath": "src\\ell\\lmp\\complex.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\complex.py",
            "file_name": "complex.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 201,
            "span_ids": [
                "complex"
            ],
            "start_line": 81,
            "end_line": 92,
            "community": null
        },
        "node_id": "lmp\\complex.py::3"
    },
    {
        "id": "lmp\\complex.py::4",
        "input_type": "file",
        "content": "def _get_messages(prompt_ret: Union[str, list[MessageOrDict]], prompt: LMP) -> list[Message]:\n    \"\"\"\n    Helper function to convert the output of an LMP into a list of Messages.\n    \"\"\"\n    if isinstance(prompt_ret, str):\n        has_system_prompt = prompt.__doc__ is not None and prompt.__doc__.strip() != \"\"\n        messages =     [Message(role=\"system\", content=[ContentBlock(text=_lstr(prompt.__doc__ ) )])] if has_system_prompt else []\n        return messages + [\n            Message(role=\"user\", content=[ContentBlock(text=prompt_ret)])\n        ]\n    else:\n        assert isinstance(\n            prompt_ret, list\n        ), \"Need to pass a list of Messages to the language model\"\n        return prompt_ret",
        "filepath": "src\\ell\\lmp\\complex.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\complex.py",
            "file_name": "complex.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 162,
            "span_ids": [
                "_get_messages"
            ],
            "start_line": 96,
            "end_line": 110,
            "community": null
        },
        "node_id": "lmp\\complex.py::4"
    },
    {
        "id": "lmp\\complex.py::5",
        "input_type": "file",
        "content": "def _client_for_model(\n    model: str,\n    client: Optional[Any] = None,\n    _name: Optional[str] = None,\n) -> Any:\n    # XXX: Move to config to centralize api keys etc.\n    if not client:\n        client, was_fallback = config.get_client_for(model)\n\n        # XXX: Wrong.\n        if not client and not was_fallback:\n            raise RuntimeError(_no_api_key_warning(model, _name, '', long=True, error=True))\n\n    if client is None:\n        raise ValueError(f\"No client found for model '{model}'. Ensure the model is registered using 'register_model' in 'config.py' or specify a client directly using the 'client' argument in the decorator or function call.\")\n    return client\n\n\ncomplex.__doc__ =\n # ... other code",
        "filepath": "src\\ell\\lmp\\complex.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\complex.py",
            "file_name": "complex.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 166,
            "span_ids": [
                "impl",
                "_client_for_model"
            ],
            "start_line": 112,
            "end_line": 327,
            "community": null
        },
        "node_id": "lmp\\complex.py::5"
    },
    {
        "id": "lmp\\simple.py::1",
        "input_type": "file",
        "content": "from functools import wraps\nfrom typing import Any, Optional\n\nfrom ell.lmp.complex import complex\n\n\ndef simple(model: str, client: Optional[Any] = None,  exempt_from_tracking=False, **api_params):\n    assert 'tools' not in api_params, \"tools are not supported in lm decorator, use multimodal decorator instead\"\n    assert 'tool_choice' not in api_params, \"tool_choice is not supported in lm decorator, use multimodal decorator instead\"\n    assert 'response_format' not in api_params or isinstance(api_params.get('response_format', None), dict), \"response_format is not supported in lm decorator, use multimodal decorator instead\"\n\n    def convert_multimodal_response_to_lstr(response):\n        return [x.content[0].text for x in response] if isinstance(response, list) else response.content[0].text\n    return complex(model, client,  exempt_from_tracking=exempt_from_tracking, **api_params, post_callback=convert_multimodal_response_to_lstr)",
        "filepath": "src\\ell\\lmp\\simple.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\simple.py",
            "file_name": "simple.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 209,
            "span_ids": [
                "simple",
                "imports"
            ],
            "start_line": 1,
            "end_line": 14,
            "community": null
        },
        "node_id": "lmp\\simple.py::1"
    },
    {
        "id": "lmp\\simple.py::2",
        "input_type": "file",
        "content": "simple.__doc__ = \"\"\"The fundamental unit of language model programming in ell.\n\n  This decorator simplifies the process of creating Language Model Programs (LMPs) \n  that return text-only outputs from language models, while supporting multimodal inputs.\n  It wraps the more complex 'complex' decorator, providing a streamlined interface for common use cases.\n\n  :param model: The name or identifier of the language model to use.\n  :type model: str\n  :param client: An optional OpenAI client instance. If not provided, a default client will be used.\n  :type client: Optional[openai.Client]\n  :param exempt_from_tracking: If True, the LMP usage won't be tracked. Default is False.\n  :type exempt_from_tracking: bool\n  :param api_params: Additional keyword arguments to pass to the underlying API call.\n  :type api_params: Any\n\n  Usage:\n  The decorated function can return either a single prompt or a list of ell.Message objects:\n\n  .. code-block:: python\n\n      @ell.simple(model=\"gpt-4\", temperature=0.7)\n      def summarize_text(text: str) -> str:\n          '''You are an expert at summarizing text.''' # System prompt\n          return f\"Please summarize the following text:\\\\n\\\\n{text}\" # User prompt\n\n\n      @ell.simple(model=\"gpt-4\", temperature=0.7)\n      def describe_image(image : PIL.Image.Image) -> List[ell.Message]:\n          '''Describe the contents of an image.''' # unused because we're returning a list of Messages\n          return [\n              # helper function for ell.Message(text=\"...\", role=\"system\")\n              ell.system(\"You are an AI trained to describe images.\"),\n              # helper function for ell.Message(content=\"...\", role=\"user\")\n              ell.user([\"Describe this image in detail.\", image]),\n          ]\n\n\n      image_description = describe_image(PIL.Image.open(\"https://example.com/image.jpg\"))\n      print(image_description) \n      # Output will be a string text-only description of the image\n\n      summary = summarize_text(\"Long text to summarize...\")\n      print(summary)\n      # Output will be a text-only summary\n\n  Notes:\n\n  - This decorator is designed for text-only model outputs, but supports multimodal inputs.\n  - It simplifies complex responses from language models to text-only format, regardless of \n    the model's capability for structured outputs, function calling, or multimodal outputs.\n  - For preserving complex model outputs (e.g., structured data, function calls, or multimodal \n    outputs), use the @ell.complex decorator instead. @ell.complex returns a Message object (role='assistant')\n  - The decorated function can return a string or a list of ell.Message objects for more \n    complex prompts, including multimodal inputs.\n  - If called with n > 1 in api_params, the wrapped LMP will return a list of strings for the n parallel outputs\n    of the model instead of just one string. Otherwise, it will return a single string.\n  - You can pass LM API parameters either in the decorator or when calling the decorated function.\n    Parameters passed during the function call will override those set in the decorator.\n\n  Example of passing LM API params:\n\n  .. code-block:: python\n\n      @ell.simple(model=\"gpt-4\", temperature=0.7)\n      def generate_story(prompt: str) -> str:\n          return f\"Write a short story based on this prompt: {prompt}\"\n\n      # Using default parameters\n      story1 = generate_story(\"A day in the life of a time traveler\")\n\n      # Overriding parameters during function call\n      story2 = generate_story(\"An AI's first day of consciousness\", api_params={\"temperature\": 0.9, \"max_tokens\": 500})\n\n  See Also:\n\n  - :func:`ell.complex`: For LMPs that preserve full structure of model responses, including multimodal outputs.\n  - :func:`ell.tool`: For defining tools that can be used within complex LMPs.\n  - :mod:`ell.studio`: For visualizing and analyzing LMP executions.\n    \"\"\"",
        "filepath": "src\\ell\\lmp\\simple.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\simple.py",
            "file_name": "simple.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 861,
            "span_ids": [
                "impl"
            ],
            "start_line": 18,
            "end_line": 96,
            "community": null
        },
        "node_id": "lmp\\simple.py::2"
    },
    {
        "id": "lmp\\tool.py::1",
        "input_type": "file",
        "content": "from functools import wraps\nimport json\nfrom typing import Any, Callable, Optional\n\nfrom pydantic import Field, create_model\nfrom pydantic.fields import FieldInfo\nfrom ell.lmp._track import _track\n# from ell.types import ToolFunction, InvocableTool, ToolParams\n# from ell.util.verbosity import compute_color, tool_usage_logger_pre\nfrom ell.configurator import config\nfrom ell.types._lstr import _lstr\nfrom ell.types.studio import LMPType\nimport inspect\n\nfrom ell.types.message import ContentBlock, InvocableTool, ToolResult, to_content_blocks",
        "filepath": "src\\ell\\lmp\\tool.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\tool.py",
            "file_name": "tool.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 127,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 15,
            "community": null
        },
        "node_id": "lmp\\tool.py::1"
    },
    {
        "id": "lmp\\tool.py::2",
        "input_type": "file",
        "content": "def tool(*, exempt_from_tracking: bool = False, **tool_kwargs):\n    def tool_decorator(fn: Callable[..., Any]) -> InvocableTool:\n        _under_fn = fn\n\n        @wraps(fn)\n        def wrapper(\n            *fn_args,\n            _invocation_origin: str = None,\n            _tool_call_id: str = None,\n            **fn_kwargs\n        ):\n            #XXX: Post release, we need to wrap all tool arguments in type primitives for tracking I guess or change that tool makes the tool function inoperable.\n            #XXX: Most people are not going to manually try and call the tool without a type primitive and if they do it will most likely be wrapped with l strs.\n\n            if config.verbose and not exempt_from_tracking:\n                pass\n                # tool_usage_logger_pre(fn, fn_args, fn_kwargs, name, color)\n\n            result = fn(*fn_args, **fn_kwargs)\n\n            _invocation_api_params = dict(tool_kwargs=tool_kwargs)\n\n            # Here you might want to add logic for tracking the tool usage\n            # Similar to how it's done in the lm decorator # Use _invocation_origin\n\n            if isinstance(result, str) and _invocation_origin:\n                result = _lstr(result,origin_trace=_invocation_origin)\n\n            #XXX: This _tool_call_id thing is a hack. Tracking should happen via params in the api\n            # So if you call wiuth a _tool_callId\n            if _tool_call_id:\n                # XXX: TODO: MOVE TRACKING CODE TO _TRACK AND OUT OF HERE AND API.\n                try:\n                    if isinstance(result, ContentBlock):\n                        content_results = [result]\n                    elif isinstance(result, list) and all(isinstance(c, ContentBlock) for c in result):\n                        content_results = result\n                    else:\n                        content_results = [ContentBlock(text=_lstr(json.dumps(result, ensure_ascii=False),origin_trace=_invocation_origin))]\n                except TypeError as e:\n                    raise TypeError(f\"Failed to convert tool use result to ContentBlock: {e}. Tools must return json serializable objects. or a list of ContentBlocks.\")\n                # XXX: Need to support images and other content types somehow. We should look for images inside of the the result and then go from there.\n                # try:\n                #     content_results = coerce_content_list(result)\n                # except ValueError as e:\n\n                # TODO: poolymorphic validation here is important (cant have tool_call or formatted_response in the result)\n                # XXX: Should we put this coercion here or in the tool call/result area.\n                for c in content_results:\n                    assert not c.tool_call, \"Tool call in tool result\"\n                    # assert not c.formatted_response, \"Formatted response in tool result\"\n                    if c.parsed:\n                        # Warning: Formatted response in tool result will be converted to text\n                        # TODO: Logging needs to produce not print.\n                        print(f\"Warning: Formatted response in tool result will be converted to text. Original: {c.parsed}\")\n                        c.text = _lstr(c.parsed.model_dump_json(),origin_trace=_invocation_origin)\n                        c.parsed = None\n                    assert not c.audio, \"Audio in tool result\"\n                return ToolResult(tool_call_id=_tool_call_id, result=content_results), _invocation_api_params, {}\n            else:\n                return result, _invocation_api_params, {}\n        # ... other code\n    # ... other code",
        "filepath": "src\\ell\\lmp\\tool.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\tool.py",
            "file_name": "tool.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 723,
            "span_ids": [
                "tool"
            ],
            "start_line": 19,
            "end_line": 79,
            "community": null
        },
        "node_id": "lmp\\tool.py::2"
    },
    {
        "id": "lmp\\tool.py::3",
        "input_type": "file",
        "content": "def tool(*, exempt_from_tracking: bool = False, **tool_kwargs):\n    def tool_decorator(fn: Callable[..., Any]) -> InvocableTool:\n        # ... other code\n\n\n        wrapper.__ell_tool_kwargs__ = tool_kwargs\n        wrapper.__ell_func__ = _under_fn\n        wrapper.__ell_type__ = LMPType.TOOL\n        wrapper.__ell_exempt_from_tracking = exempt_from_tracking\n\n        # Construct the pydantic mdoel for the _under_fn's function signature parameters.\n        # 1. Get the function signature.\n\n        sig = inspect.signature(fn)\n\n        # 2. Create a dictionary of field definitions for the Pydantic model\n        fields = {}\n        for param_name, param in sig.parameters.items():\n            # Skip *args and **kwargs\n            if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):\n                continue\n\n            # Determine the type annotation\n            if param.annotation == inspect.Parameter.empty:\n                raise ValueError(f\"Parameter {param_name} has no type annotation, and cannot be converted into a tool schema for OpenAI and other provisders. Should OpenAI produce a string or an integer, etc, for this parameter?\")\n            annotation = param.annotation\n\n            # Determine the default value\n            default = param.default\n\n            # Check if the parameter has a Field with description\n            if isinstance(param.default, FieldInfo):\n                field = param.default\n                fields[param_name] = (annotation, field)\n            elif param.default != inspect.Parameter.empty:\n                fields[param_name] = (annotation, param.default)\n            else:\n                # If no default value, use Field without default\n                fields[param_name] = (annotation, Field(...))\n\n        # 3. Create the Pydantic model\n        model_name = f\"{fn.__name__}\"\n        ParamsModel = create_model(model_name, **fields)\n\n        # Attach the Pydantic model to the wrapper function\n        wrapper.__ell_params_model__ = ParamsModel\n\n        # handle tracking last.\n        if exempt_from_tracking:\n            ret = wrapper\n        else:\n            ret=  _track(wrapper)\n\n        # Helper function to get the Pydantic model for the tool\n        def get_params_model():\n            return wrapper.__ell_params_model__\n\n        # Attach the helper function to the wrapper\n        wrapper.get_params_model = get_params_model\n        ret.get_params_model = get_params_model\n        return ret\n\n    return tool_decorator",
        "filepath": "src\\ell\\lmp\\tool.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\tool.py",
            "file_name": "tool.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 510,
            "span_ids": [
                "tool"
            ],
            "start_line": 82,
            "end_line": 139,
            "community": null
        },
        "node_id": "lmp\\tool.py::3"
    },
    {
        "id": "lmp\\tool.py::4",
        "input_type": "file",
        "content": "tool.__doc__ = \"\"\"Defines a tool for use in language model programs (LMPs) that support tool use.\n\nThis decorator wraps a function, adding metadata and handling for tool invocations.\nIt automatically extracts the tool's description and parameters from the function's\ndocstring and type annotations, creating a structured representation for LMs to use.\n\n:param exempt_from_tracking: If True, the tool usage won't be tracked. Default is False.\n:type exempt_from_tracking: bool\n:param tool_kwargs: Additional keyword arguments for tool configuration.\n:return: A wrapped version of the original function, usable as a tool by LMs.\n:rtype: Callable\n\nRequirements:\n\n- Function must have fully typed arguments (Pydantic-serializable).\n- Return value must be one of: str, JSON-serializable object, Pydantic model, or List[ContentBlock].\n- All parameters must have type annotations.\n- Complex types should be Pydantic models.\n- Function should have a descriptive docstring.\n- Can only be used in LMPs with @ell.complex decorators\n\nFunctionality:\n\n1. Metadata Extraction:\n    - Uses function docstring as tool description.\n    - Extracts parameter info from type annotations and docstring.\n    - Creates a Pydantic model for parameter validation and schema generation.\n\n2. Integration with LMs:\n    - Can be passed to @ell.complex decorators.\n    - Provides structured tool information to LMs.\n\n3. Invocation Handling:\n    - Manages tracking, logging, and result processing.\n    - Wraps results in appropriate types (e.g., _lstr) for tracking.\n\nUsage Modes:\n\n1. Normal Function Call:\n    - Behaves like a regular Python function.\n    - Example: result = my_tool(arg1=\"value\", arg2=123)\n\n2. LMP Tool Call:\n    - Used within LMPs or with explicit _tool_call_id.\n    - Returns a ToolResult object.\n    - Example: result = my_tool(arg1=\"value\", arg2=123, _tool_call_id=\"unique_id\")\n\nResult Coercion:\n\n- String \u2192 ContentBlock(text=result)\n- Pydantic BaseModel \u2192 ContentBlock(parsed=result)\n- List[ContentBlock] \u2192 Used as-is\n- Other types \u2192 ContentBlock(text=json.dumps(result))\n\nExample::\n\n    @ell.tool()\n    def create_claim_draft(\n        claim_details: str,\n        claim_type: str,\n        claim_amount: float,\n        claim_date: str = Field(description=\"Date format: YYYY-MM-DD\")\n    ) -> str:\n        '''Create a claim draft. Returns the created claim ID.'''\n        return \"12345\"\n\n    # For use in a complex LMP:\n    @ell.complex(model=\"gpt-4\", tools=[create_claim_draft], temperature=0.1)\n    def insurance_chatbot(message_history: List[Message]) -> List[Message]:\n        # Chatbot implementation...\n\n    x = insurance_chatbot([\n        ell.user(\"I crashed my car into a tree.\"),\n        ell.assistant(\"I'm sorry to hear that. Can you provide more details?\"),\n        ell.user(\"The car is totaled and I need to file a claim. Happened on 2024-08-01. total value is like $5000\")\n    ]) \n    print(x)\n    '''ell.Message(content=[\n        ContentBlock(tool_call(\n            tool_call_id=\"asdas4e\",\n            tool_fn=create_claim_draft,\n            input=create_claim_draftParams({\n                claim_details=\"The car is totaled and I need to file a claim. Happened on 2024-08-01. total value is like $5000\",\n                claim_type=\"car\",\n                claim_amount=5000,\n                claim_date=\"2024-08-01\"\n            })\n        ))\n    ], role='assistant')'''\n    \n    if x.tool_calls:\n        next_user_message = response_message.call_tools_and_collect_as_message()\n        # This actually calls create_claim_draft\n        print(next_user_message)\n        '''\n        ell.Message(content=[\n            ContentBlock(tool_result=ToolResult(\n                tool_call_id=\"asdas4e\",\n                result=[ContentBlock(text=\"12345\")]\n            ))\n        ], role='user')\n        '''\n        y = insurance_chatbot(message_history + [x, next_user_message])\n        print(y)\n        '''\n        ell.Message(\"I've filed that for you!\", role='assistant')\n        '''\n\nNote:\n- Tools are integrated into LMP calls via the 'tools' parameter in @ell.complex.\n- LMs receive structured tool information, enabling understanding and usage within the conversation context.\n    \"\"\"",
        "filepath": "src\\ell\\lmp\\tool.py",
        "metadata": {
            "file_path": "src\\ell\\lmp\\tool.py",
            "file_name": "tool.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 955,
            "span_ids": [
                "impl"
            ],
            "start_line": 142,
            "end_line": 253,
            "community": null
        },
        "node_id": "lmp\\tool.py::4"
    },
    {
        "id": "models\\__init__.py::1",
        "input_type": "file",
        "content": "\"\"\"\nAttempts to registeres model names with their respective API client bindings. This allows for the creation of a unified interface for interacting with different LLM providers.\n\nFor example, to register an OpenAI model:\n@ell.simple(model='gpt-4o-mini') -> @ell.simple(model='gpt-4o-mini', client=openai.OpenAI())\n\n\"\"\"\n\nimport ell.models.openai\nimport ell.models.anthropic\nimport ell.models.ollama\nimport ell.models.groq\nimport ell.models.bedrock",
        "filepath": "src\\ell\\models\\__init__.py",
        "metadata": {
            "file_path": "src\\ell\\models\\__init__.py",
            "file_name": "__init__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 105,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 13,
            "community": null
        },
        "node_id": "models\\__init__.py::1"
    },
    {
        "id": "models\\anthropic.py::1",
        "input_type": "file",
        "content": "from ell.configurator import config\nimport logging\n\nlogger = logging.getLogger(__name__)\n\n\ntry:\n    import anthropic\n\n    def register(client: anthropic.Anthropic):\n        \"\"\"\n        Register Anthropic models with the provided client.\n\n        This function takes an Anthropic client and registers various Anthropic models\n        with the global configuration. It allows the system to use these models\n        for different AI tasks.\n\n        Args:\n            client (anthropic.Anthropic): An instance of the Anthropic client to be used\n                                          for model registration.\n\n        Note:\n            The function doesn't return anything but updates the global\n            configuration with the registered models.\n        \"\"\"\n        model_data = [\n            ('claude-3-opus-20240229', 'anthropic'),\n            ('claude-3-sonnet-20240229', 'anthropic'),\n            ('claude-3-haiku-20240307', 'anthropic'),\n            ('claude-3-5-sonnet-20240620', 'anthropic'),\n        ]\n        for model_id, owned_by in model_data:\n            config.register_model(model_id, client)\n\n    try:\n        default_client = anthropic.Anthropic()\n        register(default_client)\n    except Exception as e:\n        # logger.warning(f\"Failed to create default Anthropic client: {e}\")\n        pass\n\n\nexcept ImportError:\n    pass",
        "filepath": "src\\ell\\models\\anthropic.py",
        "metadata": {
            "file_path": "src\\ell\\models\\anthropic.py",
            "file_name": "anthropic.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 284,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 44,
            "community": null
        },
        "node_id": "models\\anthropic.py::1"
    },
    {
        "id": "models\\bedrock.py::1",
        "input_type": "file",
        "content": "from typing import Any\nfrom ell.configurator import config\nimport logging\n\nlogger = logging.getLogger(__name__)\n\n\ndef register(client: Any):\n    \"\"\"\n    Register Bedrock models with the provided client.\n\n    This function takes an boto3 client and registers various Bedrock models\n    with the global configuration. It allows the system to use these models\n    for different AI tasks.\n\n    Args:\n        client (boto3.client): An instance of the bedrock client to be used\n                                        for model registration.\n\n    Note:\n        The function doesn't return anything but updates the global\n        configuration with the registered models.\n    \"\"\"\n    model_data = [\n        ('anthropic.claude-3-opus-20240229-v1:0', 'bedrock'),\n        ('anthropic.claude-3-sonnet-20240229-v1:0', 'bedrock'),\n        ('anthropic.claude-3-haiku-20240307-v1:0', 'bedrock'),\n        ('anthropic.claude-3-5-sonnet-20240620-v1:0', 'bedrock'),\n\n        ('mistral.mistral-7b-instruct-v0:2', 'bedrock'),\n        ('mistral.mixtral-8x7b-instruct-v0:1', 'bedrock'),\n        ('mistral.mistral-large-2402-v1:0', 'bedrock'),\n        ('mistral.mistral-small-2402-v1:0', 'bedrock'),\n\n\n        ('ai21.jamba-instruct-v1:0','bedrock'),\n        ('ai21.j2-ultra-v1', 'bedrock'),\n        ('ai21.j2-mid-v1', 'bedrock'),\n\n        ('amazon.titan-embed-text-v1', 'bedrock'),\n        ('amazon.titan-text-lite-v1', 'bedrock'),\n        ('amazon.titan-text-express-v1', 'bedrock'),\n        ('amazon.titan-image-generator-v2:0', 'bedrock'),\n        ('amazon.titan-image-generator-v1', 'bedrock'),\n\n        ('cohere.command-r-plus-v1:0', 'bedrock'),\n        ('cohere.command-r-v1:0', 'bedrock'),\n        ('cohere.embed-english-v3', 'bedrock'),\n        ('cohere.embed-multilingual-v3', 'bedrock'),\n        ('cohere.command-text-v14', 'bedrock'),\n\n        ('meta.llama3-8b-instruct-v1:0', 'bedrock'),\n        ('meta.llama3-70b-instruct-v1:0', 'bedrock'),\n        ('meta.llama2-13b-chat-v1', 'bedrock'),\n        ('meta.llama2-70b-chat-v1', 'bedrock'),\n        ('meta.llama2-13b-v1', 'bedrock'),\n\n    ]\n\n    for model_id, owned_by in model_data:\n        config.register_model(name=model_id, default_client=client, supports_streaming=True)\n\ndefault_client = None\ntry:\n\n    import boto3\n    default_client = boto3.client('bedrock-runtime')\nexcept Exception as e:\n    pass\n\nregister(default_client)",
        "filepath": "src\\ell\\models\\bedrock.py",
        "metadata": {
            "file_path": "src\\ell\\models\\bedrock.py",
            "file_name": "bedrock.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 672,
            "span_ids": [
                "impl:3",
                "imports",
                "register"
            ],
            "start_line": 1,
            "end_line": 72,
            "community": null
        },
        "node_id": "models\\bedrock.py::1"
    },
    {
        "id": "models\\groq.py::1",
        "input_type": "file",
        "content": "from typing import Optional\nfrom ell.configurator import config\n\ntry:\n    from groq import Groq\n    def register(client: Optional[Groq] = None, **client_kwargs):\n        if client is None:\n            client = Groq(**client_kwargs)\n        for model in client.models.list().data:\n            config.register_model(model.id, default_client=client, supports_streaming=True)\nexcept ImportError:\n    pass",
        "filepath": "src\\ell\\models\\groq.py",
        "metadata": {
            "file_path": "src\\ell\\models\\groq.py",
            "file_name": "groq.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 86,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 12,
            "community": null
        },
        "node_id": "models\\groq.py::1"
    },
    {
        "id": "models\\ollama.py::1",
        "input_type": "file",
        "content": "from ell.configurator import config\nimport openai\nimport requests\nimport logging\n\n#XXX: May be deprecated soon because of the new provider framework.\nlogger = logging.getLogger(__name__)\nclient = None\n\ndef register(base_url):\n    \"\"\"\n    Registers Ollama models with the provided base URL.\n\n    This function sets up the Ollama client with the given base URL and\n    fetches available models from the Ollama API. It then registers these\n    models with the global configuration, allowing them to be used within\n    the ell framework.\n\n    Args:\n        base_url (str): The base URL of the Ollama API endpoint.\n\n    Note:\n        This function updates the global client and configuration.\n        It logs any errors encountered during the process.\n    \"\"\"\n    global client\n    client = openai.Client(base_url=base_url)\n\n    try:\n        response = requests.get(f\"{base_url}/../api/tags\")\n        response.raise_for_status()\n        models = response.json().get(\"models\", [])\n\n        for model in models:\n            config.register_model(model[\"name\"], client)\n    except requests.RequestException as e:\n        logger.error(f\"Failed to fetch models from {base_url}: {e}\")\n    except Exception as e:\n        logger.error(f\"An error occurred: {e}\")",
        "filepath": "src\\ell\\models\\ollama.py",
        "metadata": {
            "file_path": "src\\ell\\models\\ollama.py",
            "file_name": "ollama.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 268,
            "span_ids": [
                "imports",
                "register"
            ],
            "start_line": 1,
            "end_line": 41,
            "community": null
        },
        "node_id": "models\\ollama.py::1"
    },
    {
        "id": "models\\openai.py::1",
        "input_type": "file",
        "content": "\"\"\"\nThis module handles the registration of OpenAI models within the ell framework.\n\nIt provides functionality to register various OpenAI models with a given OpenAI client,\nmaking them available for use throughout the system. The module also sets up a default\nclient behavior for unregistered models.\n\nKey features:\n1. Registration of specific OpenAI models with their respective types (system, openai, openai-internal).\n2. Utilization of a default OpenAI client for any unregistered models,\n\nThe default client behavior ensures that even if a specific model is not explicitly\nregistered, the system can still attempt to use it with the default OpenAI client.\nThis fallback mechanism provides flexibility in model usage while maintaining a\nstructured approach to model registration.\n\nNote: The actual model availability may depend on your OpenAI account's access and the\ncurrent offerings from OpenAI.\n\nAdditionally, due to the registration of default mdoels, the OpenAI client may be used for\nanthropic, cohere, groq, etc. models if their clients are not registered or fail\nto register due to an error (lack of API keys, rate limits, etc.)\n\"\"\"\n\nfrom ell.configurator import config\nimport openai\n\nimport logging\nimport colorama\n\nlogger = logging.getLogger(__name__)",
        "filepath": "src\\ell\\models\\openai.py",
        "metadata": {
            "file_path": "src\\ell\\models\\openai.py",
            "file_name": "openai.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 258,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 31,
            "community": null
        },
        "node_id": "models\\openai.py::1"
    },
    {
        "id": "models\\openai.py::2",
        "input_type": "file",
        "content": "def register(client: openai.Client):\n    \"\"\"\n    Register OpenAI models with the provided client.\n\n    This function takes an OpenAI client and registers various OpenAI models\n    with the global configuration. It allows the system to use these models\n    for different AI tasks.\n\n    Args:\n        client (openai.Client): An instance of the OpenAI client to be used\n                                for model registration.\n\n    Note:\n        The function doesn't return anything but updates the global\n        configuration with the registered models.\n    \"\"\"\n    #XXX: Deprecation in 0.1.0\n    standard_models = [\n        'gpt-4-1106-preview',\n        'gpt-4-32k-0314',\n        'text-embedding-3-large',\n        'gpt-4-0125-preview',\n        'babbage-002',\n        'gpt-4-turbo-preview',\n        'gpt-4o',\n        'gpt-4o-2024-05-13',\n        'gpt-4o-mini-2024-07-18',\n        'gpt-4o-mini',\n        'gpt-4o-2024-08-06',\n        'gpt-3.5-turbo-0301',\n        'gpt-3.5-turbo-0613',\n        'tts-1',\n        'gpt-3.5-turbo',\n        'gpt-3.5-turbo-16k',\n        'davinci-002',\n        'gpt-3.5-turbo-16k-0613',\n        'gpt-4-turbo-2024-04-09',\n        'gpt-3.5-turbo-0125',\n        'gpt-4-turbo',\n        'gpt-3.5-turbo-1106',\n        'gpt-3.5-turbo-instruct-0914',\n        'gpt-3.5-turbo-instruct',\n        'gpt-4-0613',\n        'gpt-4',\n        'gpt-4-0314',\n        'gpt-4o-audio-preview',\n        'gpt-4o-realtime',\n    ]\n    for model_id in standard_models:\n        config.register_model(model_id, client)\n\n    #XXX: Deprecation in 0.1.0\n    config.register_model('o1-preview', client, supports_streaming=False)\n    config.register_model('o1-mini', client, supports_streaming=False)\n\ndefault_client = None\ntry:\n    default_client = openai.Client()\nexcept openai.OpenAIError as e:\n    pass\n\nregister(default_client)\nconfig.default_client = default_client",
        "filepath": "src\\ell\\models\\openai.py",
        "metadata": {
            "file_path": "src\\ell\\models\\openai.py",
            "file_name": "openai.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 580,
            "span_ids": [
                "impl:3",
                "register"
            ],
            "start_line": 33,
            "end_line": 95,
            "community": null
        },
        "node_id": "models\\openai.py::2"
    },
    {
        "id": "ell\\provider.py::1",
        "input_type": "file",
        "content": "from abc import ABC, abstractmethod\nfrom collections import defaultdict\nfrom functools import lru_cache\nimport inspect\nfrom types import MappingProxyType\nfrom typing import (\n    Any,\n    Callable,\n    Dict,\n    FrozenSet,\n    List,\n    Optional,\n    Set,\n    Tuple,\n    Type,\n    TypedDict,\n    Union,\n)\n\nfrom pydantic import BaseModel, ConfigDict, Field\nfrom ell.types import Message, ContentBlock, ToolCall\nfrom ell.types._lstr import _lstr\nimport json\nfrom dataclasses import dataclass\nfrom ell.types.message import LMP",
        "filepath": "src\\ell\\provider.py",
        "metadata": {
            "file_path": "src\\ell\\provider.py",
            "file_name": "provider.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 122,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 25,
            "community": null
        },
        "node_id": "ell\\provider.py::1"
    },
    {
        "id": "ell\\provider.py::2",
        "input_type": "file",
        "content": "# XXX: Might leave this internal to providers so that the complex code is simpler &\n# we can literally jsut call provider.call like any openai fn.\nclass EllCallParams(BaseModel):\n    model: str = Field(..., description=\"Model identifier\")\n    messages: List[Message] = Field(..., description=\"Conversation context\")\n    client: Any = Field(..., description=\"API client\")\n    tools: List[LMP] = Field(default_factory=list, description=\"Available tools\")\n    api_params: Dict[str, Any] = Field(\n        default_factory=dict, description=\"API parameters\"\n    )\n\n    model_config = ConfigDict(arbitrary_types_allowed=True)\n\n    def get_tool_by_name(self, name: str) -> Optional[LMP]:\n        \"\"\"Get a tool by name.\"\"\"\n        return next(\n            (tool for tool in (self.tools or [])  if tool.__name__ == name), None\n        )\n\n\nMetadata = Dict[str, Any]",
        "filepath": "src\\ell\\provider.py",
        "metadata": {
            "file_path": "src\\ell\\provider.py",
            "file_name": "provider.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 196,
            "span_ids": [
                "EllCallParams",
                "impl",
                "imports",
                "EllCallParams.get_tool_by_name"
            ],
            "start_line": 28,
            "end_line": 48,
            "community": null
        },
        "node_id": "ell\\provider.py::2"
    },
    {
        "id": "ell\\provider.py::3",
        "input_type": "file",
        "content": "# XXX: Needs a better name.\nclass Provider(ABC):\n    \"\"\"\n    Abstract base class for all providers. Providers are API interfaces to language models, not necessarily API providers.\n    For example, the OpenAI provider is an API interface to OpenAI's API but also to Ollama and Azure OpenAI.\n    In Ell. We hate abstractions. The only reason this exists is to force implementers to implement their own provider correctly -_-.\n    \"\"\"\n    dangerous_disable_validation = False\n\n    ################################\n    ### API PARAMETERS #############\n    ################################\n    @abstractmethod\n    def provider_call_function(\n        self, client: Any, api_call_params: Optional[Dict[str, Any]] = None\n    ) -> Callable[..., Any]:\n        \"\"\"\n        Implement this method to return the function that makes the API call to the language model.\n        For example, if you're implementing the OpenAI provider, you would return the function that makes the API call to OpenAI's API.\n        \"\"\"\n        return NotImplemented\n\n    def disallowed_api_params(self) -> FrozenSet[str]:\n        \"\"\"\n        Returns a list of disallowed call params that ell will override.\n        \"\"\"\n        return frozenset({\"messages\", \"tools\", \"model\", \"stream\", \"stream_options\"})\n\n    def available_api_params(self, client: Any, api_params: Optional[Dict[str, Any]] = None):\n        params = _call_params(self.provider_call_function(client, api_params))\n        return frozenset(params.keys()) - self.disallowed_api_params()\n\n    ################################\n    ### TRANSLATION ###############\n    ################################\n    @abstractmethod\n    def translate_to_provider(self, ell_call: EllCallParams) -> Dict[str, Any]:\n        \"\"\"Converts an ell call to provider call params!\"\"\"\n        return NotImplemented\n\n    @abstractmethod\n    def translate_from_provider(\n        self,\n        provider_response: Any,\n        ell_call: EllCallParams,\n        provider_call_params: Dict[str, Any],\n        origin_id: Optional[str] = None,\n        logger: Optional[Callable[..., None]] = None,\n    ) -> Tuple[List[Message], Metadata]:\n        \"\"\"Converts provider responses to universal format. with metadata\"\"\"\n        return NotImplemented\n\n    ################################\n    ### CALL MODEL ################\n    ################################\n    # Be careful to override this method in your provider.\r",
        "filepath": "src\\ell\\provider.py",
        "metadata": {
            "file_path": "src\\ell\\provider.py",
            "file_name": "provider.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 479,
            "span_ids": [
                "Provider.available_api_params",
                "Provider",
                "impl",
                "Provider.provider_call_function",
                "Provider.translate_to_provider",
                "Provider.translate_from_provider",
                "Provider.disallowed_api_params"
            ],
            "start_line": 50,
            "end_line": 105,
            "community": null
        },
        "node_id": "ell\\provider.py::3"
    },
    {
        "id": "ell\\provider.py::4",
        "input_type": "file",
        "content": "class Provider(ABC):\n    def call(\n        self,\n        #XXX: In future refactors, we can fully enumerate the args and make ell_call's internal to the _provider implementer interface.\n        # This gives us a litellm style interface for free.\n        ell_call: EllCallParams,\n        origin_id: Optional[str] = None,\n        logger: Optional[Any] = None,\n    ) -> Tuple[List[Message], Dict[str, Any], Metadata]:\n        # Automatic validation of params\n        assert (\n            not set(ell_call.api_params.keys()).intersection(self.disallowed_api_params()) \n        ), f\"Disallowed api parameters: {ell_call.api_params}\"\n\n        final_api_call_params = self.translate_to_provider(ell_call)\n\n        call = self.provider_call_function(ell_call.client, final_api_call_params)\n        assert self.dangerous_disable_validation or _validate_provider_call_params(final_api_call_params, call)\n\n\n        provider_resp = call(**final_api_call_params)\n\n        messages, metadata = self.translate_from_provider(\n            provider_resp, ell_call, final_api_call_params, origin_id, logger\n        )\n        assert \"choices\" not in metadata, \"choices should be in the metadata.\"\n        assert self.dangerous_disable_validation or _validate_messages_are_tracked(messages, origin_id)\n\n        return messages, final_api_call_params, metadata",
        "filepath": "src\\ell\\provider.py",
        "metadata": {
            "file_path": "src\\ell\\provider.py",
            "file_name": "provider.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 279,
            "span_ids": [
                "Provider.call"
            ],
            "start_line": 106,
            "end_line": 133,
            "community": null
        },
        "node_id": "ell\\provider.py::4"
    },
    {
        "id": "ell\\provider.py::5",
        "input_type": "file",
        "content": "# handhold the the implementer, in production mode we can turn these off for speed.\n@lru_cache(maxsize=None)\ndef _call_params(call: Callable[..., Any]) -> MappingProxyType[str, inspect.Parameter]:\n    return inspect.signature(call).parameters\n\n\ndef _validate_provider_call_params(\n    api_call_params: Dict[str, Any], call: Callable[..., Any]\n):\n    provider_call_params = _call_params(call)\n\n    required_params = {\n        name: param\n        for name, param in provider_call_params.items()\n        if param.default == param.empty and param.kind != param.VAR_KEYWORD\n    }\n\n    for param_name in required_params:\n        assert (\n            param_name in api_call_params\n        ), f\"Provider implementation error: Required parameter '{param_name}' is missing in the converted call parameters converted from ell call.\"\n\n    for param_name, param_value in api_call_params.items():\n        assert (\n            param_name in provider_call_params\n        ), f\"Provider implementation error: Unexpected parameter '{param_name}' in the converted call parameters.\"\n\n    return True",
        "filepath": "src\\ell\\provider.py",
        "metadata": {
            "file_path": "src\\ell\\provider.py",
            "file_name": "provider.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 219,
            "span_ids": [
                "_validate_provider_call_params",
                "Provider.call",
                "_call_params"
            ],
            "start_line": 136,
            "end_line": 163,
            "community": null
        },
        "node_id": "ell\\provider.py::5"
    },
    {
        "id": "ell\\provider.py::6",
        "input_type": "file",
        "content": "def _validate_messages_are_tracked(\n    messages: List[Message], origin_id: Optional[str] = None\n):\n    if origin_id is None:\n        return\n\n    for message in messages:\n        assert isinstance(\n            message.text, _lstr\n        ), f\"Provider implementation error: Message text should be an instance of _lstr, got {type(message.text)}\"\n        assert (\n            origin_id in message.text.__origin_trace__\n        ), f\"Provider implementation error: Message origin_id {message.text.__origin_trace__} does not match the provided origin_id {origin_id}\"\n    return True",
        "filepath": "src\\ell\\provider.py",
        "metadata": {
            "file_path": "src\\ell\\provider.py",
            "file_name": "provider.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 124,
            "span_ids": [
                "_validate_messages_are_tracked"
            ],
            "start_line": 165,
            "end_line": 179,
            "community": null
        },
        "node_id": "ell\\provider.py::6"
    },
    {
        "id": "providers\\__init__.py::1",
        "input_type": "file",
        "content": "import ell.providers.openai\nimport ell.providers.groq\nimport ell.providers.anthropic\nimport ell.providers.bedrock\n# import ell.providers.mistral\n# import ell.providers.cohere\n# import ell.providers.gemini\n# import ell.providers.elevenlabs\n# import ell.providers.replicate\n# import ell.providers.huggingface",
        "filepath": "src\\ell\\providers\\__init__.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\__init__.py",
            "file_name": "__init__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 73,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 10,
            "community": null
        },
        "node_id": "providers\\__init__.py::1"
    },
    {
        "id": "providers\\anthropic.py::1",
        "input_type": "file",
        "content": "from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Type, Union, cast\nfrom ell.provider import  EllCallParams, Metadata, Provider\nfrom ell.types import Message, ContentBlock, ToolCall, ImageContent\n\nfrom ell.types._lstr import _lstr\nfrom ell.types.message import LMP\nfrom ell.configurator import register_provider\nfrom ell.util.serialization import serialize_image\nimport base64\nfrom io import BytesIO\nimport json\nimport requests\nfrom PIL import Image as PILImage\n\ntry:\n    import anthropic\n    from anthropic import Anthropic\n    from anthropic.types import Message as AnthropicMessage, MessageParam, RawMessageStreamEvent\n    from anthropic.types.message_create_params import MessageCreateParamsStreaming\n    from anthropic._streaming import Stream\n\n    class AnthropicProvider(Provider):\n        dangerous_disable_validation = True\n\n        def provider_call_function(self, client : Anthropic, api_call_params : Optional[Dict[str, Any]] = None) -> Callable[..., Any]:\n            return client.messages.create\n\n        def translate_to_provider(self, ell_call : EllCallParams):\n            final_call_params = cast(MessageCreateParamsStreaming, ell_call.api_params.copy())\n            # XXX: Helper, but should be depreicated due to ssot\n            assert final_call_params.get(\"max_tokens\") is not None, f\"max_tokens is required for anthropic calls, pass it to the @ell.simple/complex decorator, e.g. @ell.simple(..., max_tokens=your_max_tokens) or pass it to the model directly as a parameter when calling your LMP: your_lmp(..., api_params=({{'max_tokens': your_max_tokens}})).\"\n\n            dirty_msgs = [\n                MessageParam(\n                    role=cast(Literal[\"user\", \"assistant\"], message.role), \n                    content=[_content_block_to_anthropic_format(c) for c in message.content]) for message in ell_call.messages]\n            role_correct_msgs   : List[MessageParam] = []\n            for msg in dirty_msgs:\n                if (not len(role_correct_msgs) or role_correct_msgs[-1]['role'] != msg['role']):\n                    role_correct_msgs.append(msg)\n                else: cast(List, role_correct_msgs[-1]['content']).extend(msg['content'])\n\n            system_message = None\n            if role_correct_msgs and role_correct_msgs[0][\"role\"] == \"system\":\n                system_message = role_correct_msgs.pop(0)\n\n            if system_message:\n                final_call_params[\"system\"] = system_message[\"content\"][0][\"text\"]\n\n\n            final_call_params['stream'] = True\n            final_call_params[\"model\"] = ell_call.model\n            final_call_params[\"messages\"] = role_correct_msgs\n\n            if ell_call.tools:\n                final_call_params[\"tools\"] = [\n                    #XXX: Cleaner with LMP's as a class.\n                    dict(\n                        name=tool.__name__,\n                        description=tool.__doc__,\n                        input_schema=tool.__ell_params_model__.model_json_schema(),\n                    )\n                    for tool in ell_call.tools\n                ]\n\n            # print(final_call_params)\n            return final_call_params\n\n        def translate_from_provider(\n            self,\n            provider_response : Union[Stream[RawMessageStreamEvent], AnthropicMessage],\n            ell_call: EllCallParams,\n            provider_call_params: Dict[str, Any],\n            origin_id: Optional[str] = None,\n            logger: Optional[Callable[..., None]] = None,\n        ) -> Tuple[List[Message], Metadata]:\n\n            usage = {}\n            tracked_results = []\n            metadata = {}\n\n            #XXX: Support n > 0\n\n            if provider_call_params.get(\"stream\", False):\n                content = []\n                current_blocks: Dict[int, Dict[str, Any]] = {}\n                message_metadata = {}\n\n                with cast(Stream[RawMessageStreamEvent], provider_response) as stream:\n                    for chunk in stream:\n                        if chunk.type == \"message_start\":\n                            message_metadata = chunk.message.model_dump()\n                            message_metadata.pop(\"content\", None)  # Remove content as we'll build it separately\n\n                        elif chunk.type == \"content_block_start\":\n                            block = chunk.content_block.model_dump()\n                            current_blocks[chunk.index] = block\n                            if block[\"type\"] == \"tool_use\":\n                                if logger: logger(f\" <tool_use: {block['name']}(\")\n                                block[\"input\"] = \"\" # force it to be a string, XXX: can implement partially parsed json later.\n                        elif chunk.type == \"content_block_delta\":\n                            if chunk.index in current_blocks:\n                                block = current_blocks[chunk.index]\n                                if (delta := chunk.delta).type == \"text_delta\":\n                                    block[\"text\"] += delta.text\n                                    if logger: logger(delta.text)\n                                if delta.type == \"input_json_delta\":\n                                    block[\"input\"] += delta.partial_json\n                                    if logger: logger(delta.partial_json)\n\n                        elif chunk.type == \"content_block_stop\":\n                            if chunk.index in current_blocks:\n                                block = current_blocks.pop(chunk.index)\n                                if block[\"type\"] == \"text\":\n                                    content.append(ContentBlock(text=_lstr(block[\"text\"],origin_trace=origin_id)))\n                                elif block[\"type\"] == \"tool_use\":\n                                    try:\n                                        matching_tool = ell_call.get_tool_by_name(block[\"name\"])\n                                        if matching_tool:\n                                            content.append(\n                                                ContentBlock(\n                                                    tool_call=ToolCall(\n                                                        tool=matching_tool,\n                                                        tool_call_id=_lstr(\n                                                            block['id'],origin_trace=origin_id\n                                                        ),\n                                                        params=json.loads(block['input']) if block['input'] else {},\n                                                    )\n                                                )\n                                            )\n                                    except json.JSONDecodeError:\n                                        if logger: logger(f\" - FAILED TO PARSE JSON\")\n                                        pass\n                                    if logger: logger(f\")>\")\n\n                        elif chunk.type == \"message_delta\":\n                            message_metadata.update(chunk.delta.model_dump())\n                            if chunk.usage:\n                                usage.update(chunk.usage.model_dump())\n\n                        elif chunk.type == \"message_stop\":\n                            tracked_results.append(Message(role=\"assistant\", content=content))\n\n                        # print(chunk)\n                metadata = message_metadata\n\n            # process metadata for ell\n            # XXX: Unify an ell metadata format for ell studio.\n            usage[\"prompt_tokens\"] = usage.get(\"input_tokens\", 0)\n            usage[\"completion_tokens\"] = usage.get(\"output_tokens\", 0)\n            usage[\"total_tokens\"] = usage['prompt_tokens'] + usage['completion_tokens']\n\n            metadata[\"usage\"] = usage\n            return tracked_results, metadata\n\n    # XXX: Make a singleton.\n    anthropic_provider = AnthropicProvider()\n    register_provider(anthropic_provider, anthropic.Anthropic)\n    register_provider(anthropic_provider, anthropic.AnthropicBedrock)\n    register_provider(anthropic_provider, anthropic.AnthropicVertex)\n\nexcept ImportError:\n    pass",
        "filepath": "src\\ell\\providers\\anthropic.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\anthropic.py",
            "file_name": "anthropic.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 1418,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 163,
            "community": null
        },
        "node_id": "providers\\anthropic.py::1"
    },
    {
        "id": "providers\\anthropic.py::2",
        "input_type": "file",
        "content": "def serialize_image_for_anthropic(img : ImageContent):\n    if img.url:\n        # Download the image from the URL\n        response = requests.get(img.url)\n        response.raise_for_status()  # Raise an exception for bad responses\n        pil_image = PILImage.open(BytesIO(response.content))\n    elif img.image:\n        pil_image = img.image\n    else:\n        raise ValueError(\"Image object has neither url nor image data.\")\n    buffer = BytesIO()\n    pil_image.save(buffer, format=\"PNG\")\n    base64_image =  base64.b64encode(buffer.getvalue()).decode()\n    return dict(\n        type=\"image\",\n        source=dict(\n            type=\"base64\",\n            media_type=\"image/png\",\n            data=base64_image\n        )\n    )",
        "filepath": "src\\ell\\providers\\anthropic.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\anthropic.py",
            "file_name": "anthropic.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 158,
            "span_ids": [
                "serialize_image_for_anthropic"
            ],
            "start_line": 165,
            "end_line": 185,
            "community": null
        },
        "node_id": "providers\\anthropic.py::2"
    },
    {
        "id": "providers\\anthropic.py::3",
        "input_type": "file",
        "content": "def _content_block_to_anthropic_format(content_block: ContentBlock):\n        if (image := content_block.image): return serialize_image_for_anthropic(image)\n        elif ((text := content_block.text) is not None): return dict(type=\"text\", text=text)\n        elif (parsed := content_block.parsed):\n            return dict(type=\"text\", text=json.dumps(parsed.model_dump(), ensure_ascii=False))\n        elif (tool_call := content_block.tool_call):\n            return dict(\n                type=\"tool_use\",\n                id=tool_call.tool_call_id,\n                name=tool_call.tool.__name__,\n                input=tool_call.params.model_dump()\n            )\n        elif (tool_result := content_block.tool_result):\n            return dict(\n                type=\"tool_result\",\n                tool_use_id=tool_result.tool_call_id,\n                content=[_content_block_to_anthropic_format(c) for c in tool_result.result]\n            )\n        else:\n            raise ValueError(\"Content block is not supported by anthropic\")",
        "filepath": "src\\ell\\providers\\anthropic.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\anthropic.py",
            "file_name": "anthropic.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 204,
            "span_ids": [
                "_content_block_to_anthropic_format"
            ],
            "start_line": 187,
            "end_line": 207,
            "community": null
        },
        "node_id": "providers\\anthropic.py::3"
    },
    {
        "id": "providers\\bedrock.py::1",
        "input_type": "file",
        "content": "from abc import ABC, abstractmethod\nfrom collections import defaultdict\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, cast\nfrom ell.provider import  EllCallParams, Metadata, Provider\nfrom ell.types import Message, ContentBlock, ToolCall, ImageContent\nfrom ell.types._lstr import _lstr\nimport json\nfrom ell.configurator import config, register_provider\nfrom ell.types.message import LMP\nfrom ell.util.serialization import serialize_image\nfrom io import BytesIO\nimport requests\nfrom PIL import Image as PILImage",
        "filepath": "src\\ell\\providers\\bedrock.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\bedrock.py",
            "file_name": "bedrock.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 118,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 13,
            "community": null
        },
        "node_id": "providers\\bedrock.py::1"
    },
    {
        "id": "providers\\bedrock.py::2",
        "input_type": "file",
        "content": "try:\n    from botocore.client import BaseClient\n    from botocore.eventstream import (EventStream)\n    class BedrockProvider(Provider):\n        dangerous_disable_validation = True\n\n        def provider_call_function(self, client : Any, api_call_params : Optional[Dict[str, Any]] = None) -> Callable[..., Any]:\n            if api_call_params and api_call_params.get(\"stream\", False):\n                api_call_params.pop('stream')\n                return client.converse_stream\n            else:\n                return client.converse\n\n        def translate_to_provider(self, ell_call : EllCallParams):\n            final_call_params = {}\n\n            if ell_call.api_params.get('api_params',{}).get('stream', False):\n                final_call_params['stream'] = ell_call.api_params.get('api_params',{}).get('stream', False)\n\n            bedrock_converse_messages = [message_to_bedrock_message_format(message) for message in ell_call.messages]\n\n            system_message = None\n            if bedrock_converse_messages and bedrock_converse_messages[0][\"role\"] == \"system\":\n                system_message = bedrock_converse_messages.pop(0)\n\n            if system_message:\n                final_call_params[\"system\"] = [{'text':system_message[\"content\"][0][\"text\"]}]\n\n            final_call_params[\"modelId\"] = ell_call.model\n            final_call_params[\"messages\"] = bedrock_converse_messages\n\n            if ell_call.tools:\n                tools = [\n                    #XXX: Cleaner with LMP's as a class.\n                    dict(\n                        toolSpec = dict(\n                            name=tool.__name__,\n                            description=tool.__doc__,\n                            inputSchema=dict(\n                                json=tool.__ell_params_model__.model_json_schema(),\n                            )\n                        )\n                    )\n                    for tool in ell_call.tools\n                ]\n                final_call_params[\"toolConfig\"] = {'tools':tools}\n\n            return final_call_params\n\n        def translate_from_provider(\n                self,\n                provider_response: Union[EventStream, Any],\n                ell_call: EllCallParams,\n                provider_call_params: Dict[str, Any],\n                origin_id: Optional[str] = None,\n                logger: Optional[Callable[..., None]] = None,\n            ) -> Tuple[List[Message], Metadata]:\n\n            usage = {}\n            metadata : Metadata = {}\n\n            metadata : Metadata = {}\n            tracked_results : List[Message] = []\n            did_stream = ell_call.api_params.get(\"api_params\", {}).get('stream')\n\n            if did_stream:\n                content = []\n                current_block: Optional[Dict[str, Any]] = {}\n                message_metadata = {}\n                for chunk in provider_response.get('stream'):\n\n                    if \"messageStart\" in chunk:\n                        current_block['content'] = ''\n                        pass\n                    elif \"contentBlockStart\" in chunk:\n                        pass\n                    elif \"contentBlockDelta\" in chunk:\n                        delta = chunk.get(\"contentBlockDelta\", {}).get(\"delta\", {})\n                        if \"text\" in delta:\n                            current_block['type'] = 'text'\n                            current_block['content'] += delta.get(\"text\")\n                            if logger:\n                                logger(delta.get(\"text\"))\n                        else:\n                            pass\n                    elif \"contentBlockStop\" in chunk:\n                        if current_block is not None:\n                            if current_block[\"type\"] == \"text\":\n                                content.append(ContentBlock(text=_lstr(content=content, origin_trace=origin_id)))\n\n                    elif \"messageStop\" in chunk:\n                        tracked_results.append(Message(role=\"assistant\", content=content))\n\n                    elif \"metadata\" in chunk:\n                        if \"usage\" in chunk[\"metadata\"]:\n                            usage[\"prompt_tokens\"] = chunk[\"metadata\"].get('usage').get(\"inputTokens\", 0)\n                            usage[\"completion_tokens\"] = chunk[\"metadata\"].get('usage').get(\"outputTokens\", 0)\n                            usage[\"total_tokens\"] = usage['prompt_tokens'] + usage['completion_tokens']\n                            message_metadata[\"usage\"] = usage\n                    else:\n                        pass\n\n\n                metadata = message_metadata\n            else:\n                # Non-streaming response processing (unchanged)\n                cbs = []\n                for content_block in provider_response.get('output', {}).get('message', {}).get('content', []):\n                    if 'text' in content_block:\n                        cbs.append(ContentBlock(text=_lstr(content_block.get('text'), origin_trace=origin_id)))\n                    elif 'toolUse' in content_block:\n                        assert ell_call.tools is not None, \"Tools were not provided to the model when calling it and yet bedrock returned a tool use.\"\n                        try:\n                            toolUse = content_block['toolUse']\n                            matching_tool = ell_call.get_tool_by_name(toolUse[\"name\"])\n                            if matching_tool:\n                                cbs.append(\n                                    ContentBlock(\n                                        tool_call=ToolCall(\n                                            tool=matching_tool,\n                                            tool_call_id=_lstr(\n                                                toolUse['toolUseId'],origin_trace=origin_id\n                                            ),\n                                            params=toolUse['input'],\n                                        )\n                                    )\n                                )\n                        except json.JSONDecodeError:\n                            if logger: logger(f\" - FAILED TO PARSE JSON\")\n                            pass\n                tracked_results.append(Message(role=\"assistant\", content=cbs))\n                if logger:\n                    logger(tracked_results[0].text)\n\n\n                # usage = call_result.response.usage.dict() if call_result.response.get('usage') else {}\n                # metadata = call_result.response.model_dump()\n                # del metadata[\"content\"]\n\n            # process metadata for ell\n            # XXX: Unify an ell metadata format for ell studio.\n            usage[\"prompt_tokens\"] = usage.get(\"inputTokens\", 0)\n            usage[\"completion_tokens\"] = usage.get(\"outputTokens\", 0)\n            usage[\"total_tokens\"] = usage['prompt_tokens'] + usage['completion_tokens']\n\n            metadata[\"usage\"] = usage\n            return tracked_results, metadata\n\n\n    # XXX: Make a singleton.\n    register_provider(BedrockProvider(), BaseClient)\nexcept ImportError:\n    pass",
        "filepath": "src\\ell\\providers\\bedrock.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\bedrock.py",
            "file_name": "bedrock.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 1226,
            "span_ids": [
                "imports"
            ],
            "start_line": 15,
            "end_line": 167,
            "community": null
        },
        "node_id": "providers\\bedrock.py::2"
    },
    {
        "id": "providers\\bedrock.py::3",
        "input_type": "file",
        "content": "def content_block_to_bedrock_format(content_block: ContentBlock) -> Dict[str, Any]:\n    if content_block.image:\n        img:ImageContent = content_block.image\n        if img.url:\n            # Download the image from the URL\n            response = requests.get(img.url)\n            response.raise_for_status()  # Raise an exception for bad responses\n            pil_image = PILImage.open(BytesIO(response.content))\n        elif img.image:\n            pil_image = img.image\n        else:\n            raise ValueError(\"Image object has neither url nor image data.\")\n        buffer = BytesIO()\n        pil_image.save(buffer, format=\"PNG\")\n        base64_image = buffer.getvalue()\n        return {\n            \"image\": {\n                \"format\": \"png\",\n                \"source\":\n                {\n                    \"bytes\": base64_image\n                }\n            }\n        }\n    elif content_block.text:\n        return {\n            \"text\": content_block.text\n        }\n    elif content_block.parsed:\n        return {\n            \"type\": \"text\",\n            \"text\": json.dumps(content_block.parsed.model_dump(), ensure_ascii=False)\n        }\n    elif content_block.tool_call:\n        return {\n            \"toolUse\": {\n                \"toolUseId\": content_block.tool_call.tool_call_id,\n                \"name\": content_block.tool_call.tool.__name__,\n                \"input\": content_block.tool_call.params.model_dump()\n            }\n        }\n    elif content_block.tool_result:\n        return {\n            \"toolResult\":{\n                \"toolUseId\": content_block.tool_result.tool_call_id,\n                \"content\": [content_block_to_bedrock_format(c) for c in content_block.tool_result.result]\n            }\n        }\n    else:\n        raise ValueError(\"Content block is not supported by bedrock\")\n\n\n\ndef message_to_bedrock_message_format(message: Message) -> Dict[str, Any]:\n\n    converse_message = {\n        \"role\": message.role,\n        \"content\": list(filter(None, [\n            content_block_to_bedrock_format(c) for c in message.content\n        ]))\n    }\n    return converse_message",
        "filepath": "src\\ell\\providers\\bedrock.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\bedrock.py",
            "file_name": "bedrock.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 420,
            "span_ids": [
                "content_block_to_bedrock_format",
                "message_to_bedrock_message_format"
            ],
            "start_line": 170,
            "end_line": 231,
            "community": null
        },
        "node_id": "providers\\bedrock.py::3"
    },
    {
        "id": "providers\\groq.py::1",
        "input_type": "file",
        "content": "\"\"\"\nGroq provider.\n\"\"\"\n\nfrom ell.providers.openai import OpenAIProvider\nfrom ell.configurator import register_provider\n\n\ntry:\n    import groq\n    class GroqProvider(OpenAIProvider):\n        dangerous_disable_validation = True\n        def translate_to_provider(self, *args, **kwargs):\n            params = super().translate_to_provider(*args, **kwargs)\n            params.pop('stream_options', None)\n            return params\n\n        def translate_from_provider(self, *args, **kwargs):\n            res, meta = super().translate_from_provider(*args, **kwargs)\n            if not meta['usage']:\n                meta['usage'] = meta['x_groq']['usage']\n            return res, meta\n    register_provider(GroqProvider(), groq.Client)\nexcept ImportError:\n    pass",
        "filepath": "src\\ell\\providers\\groq.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\groq.py",
            "file_name": "groq.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 162,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 27,
            "community": null
        },
        "node_id": "providers\\groq.py::1"
    },
    {
        "id": "providers\\openai.py::1",
        "input_type": "file",
        "content": "from abc import ABC, abstractmethod\nfrom collections import defaultdict\nfrom typing import Any, Callable, Dict, List, Optional, Tuple, Type, Union, cast\n\nfrom pydantic import BaseModel\nfrom ell.provider import  EllCallParams, Metadata, Provider\nfrom ell.types import Message, ContentBlock, ToolCall\nfrom ell.types._lstr import _lstr\nimport json\nfrom ell.configurator import _Model, config, register_provider\nfrom ell.types.message import LMP\nfrom ell.util.serialization import serialize_image\n\ntry:\n    # XXX: Could genericize.\n    import openai\n    from openai._streaming import Stream\n    from openai.types.chat import ChatCompletion, ParsedChatCompletion, ChatCompletionChunk, ChatCompletionMessageParam\n\n    class OpenAIProvider(Provider):\n        dangerous_disable_validation = True\n\n        def provider_call_function(self, client : openai.Client, api_call_params : Optional[Dict[str, Any]] = None) -> Callable[..., Any]:\n            if api_call_params and (isinstance(fmt := api_call_params.get(\"response_format\"), type)) and issubclass(fmt, BaseModel):\n                return client.beta.chat.completions.parse\n            else:\n                return client.chat.completions.create\n\n        def translate_to_provider(self, ell_call : EllCallParams) -> Dict[str, Any]:\n            final_call_params = ell_call.api_params.copy()\n            final_call_params[\"model\"] = ell_call.model\n            # Stream by default for verbose logging.\n            final_call_params[\"stream\"] = True\n            final_call_params[\"stream_options\"] = {\"include_usage\": True}\n\n            # XXX: Deprecation of config.registry.supports_streaming when streaming is implemented.\n            if ell_call.tools or final_call_params.get(\"response_format\") or (regisered_model := config.registry.get(ell_call.model, None)) and regisered_model.supports_streaming is False:\n                final_call_params.pop(\"stream\", None)\n                final_call_params.pop(\"stream_options\", None)\n            if ell_call.tools:\n                final_call_params.update(\n                    tool_choice=final_call_params.get(\"tool_choice\", \"auto\"),\n                    tools=[  \n                        dict(\n                            type=\"function\",\n                            function=dict(\n                                name=tool.__name__,\n                                description=tool.__doc__,\n                                parameters=tool.__ell_params_model__.model_json_schema(),  #type: ignore\n                            )\n                        ) for tool in ell_call.tools\n                    ]\n                )\n            # messages\n            openai_messages : List[ChatCompletionMessageParam] = []\n            for message in ell_call.messages:\n                if (tool_calls := message.tool_calls):\n                    assert message.role == \"assistant\", \"Tool calls must be from the assistant.\"\n                    assert all(t.tool_call_id for t in tool_calls), \"Tool calls must have tool call ids.\"\n                    openai_messages.append(dict(\n                        tool_calls=[\n                            dict(\n                                id=cast(str, tool_call.tool_call_id),\n                                type=\"function\",\n                                function=dict(\n                                    name=tool_call.tool.__name__,\n                                    arguments=json.dumps(tool_call.params.model_dump(), ensure_ascii=False)\n                                )\n                            ) for tool_call in tool_calls ],\n                        role=\"assistant\",\n                        content=None,\n                    ))\n                elif (tool_results := message.tool_results):\n                    for tool_result in tool_results:\n                        assert all(cb.type == \"text\" for cb in tool_result.result), \"Tool result does not match expected content blocks.\"\n                        openai_messages.append(dict(\n                            role=\"tool\",\n                            tool_call_id=tool_result.tool_call_id,\n                            content=tool_result.text_only, \n                        ))\n                else:\n                    openai_messages.append(cast(ChatCompletionMessageParam, dict(\n                        role=message.role,\n                        content=[_content_block_to_openai_format(c) for c in message.content] \n                             if message.role != \"system\" \n                             else message.text_only\n                    )))\n\n            final_call_params[\"messages\"] = openai_messages\n\n            return final_call_params\n\n        def translate_from_provider(\n            self,\n            provider_response: Union[\n                ChatCompletion, \n                ParsedChatCompletion,\n                Stream[ChatCompletionChunk], Any],\n            ell_call: EllCallParams,\n            provider_call_params: Dict[str, Any],\n            origin_id: Optional[str] = None,\n            logger: Optional[Callable[..., None]] = None,\n        ) -> Tuple[List[Message], Metadata]:\n\n            metadata : Metadata = {}\n            messages : List[Message] = []\n            did_stream = provider_call_params.get(\"stream\", False)\n\n\n            if did_stream:\n                stream = cast(Stream[ChatCompletionChunk], provider_response)\n                message_streams = defaultdict(list)\n                role : Optional[str] = None\n                for chunk in stream:\n                    metadata.update(chunk.model_dump(exclude={\"choices\"}))\n\n                    for chat_compl_chunk in chunk.choices:\n                        message_streams[chat_compl_chunk.index].append(chat_compl_chunk)\n                        delta = chat_compl_chunk.delta\n                        role = role or delta.role\n                        if  chat_compl_chunk.index == 0 and logger:\n                            logger(delta.content, is_refusal=hasattr(delta, \"refusal\") and delta.refusal)\n                for _, message_stream in sorted(message_streams.items(), key=lambda x: x[0]):\n                    text = \"\".join((choice.delta.content or \"\") for choice in message_stream)\n                    messages.append(\n                        Message(role=role, \n                                content=_lstr(content=text,origin_trace=origin_id)))\n                    #XXX: Support streaming other types.\n            else:\n                chat_completion = cast(Union[ChatCompletion, ParsedChatCompletion], provider_response)\n                metadata = chat_completion.model_dump(exclude={\"choices\"})\n                for oai_choice in chat_completion.choices:\n                    role = oai_choice.message.role\n                    content_blocks = []\n                    if (hasattr(message := oai_choice.message, \"refusal\") and (refusal := message.refusal)):\n                        raise ValueError(refusal)\n                    if hasattr(message, \"parsed\"):\n                        if (parsed := message.parsed):\n                            content_blocks.append(ContentBlock(parsed=parsed)) #XXX: Origin tracing\n                            if logger: logger(parsed.model_dump_json())\n                    else:\n                        if (content := message.content):\n                            content_blocks.append(\n                                ContentBlock(\n                                    text=_lstr(content=content,origin_trace=origin_id)))\n                            if logger: logger(content)\n                        if (tool_calls := message.tool_calls):\n                            for tool_call in tool_calls:\n                                matching_tool = ell_call.get_tool_by_name(tool_call.function.name)\n                                assert matching_tool, \"Model called tool not found in provided toolset.\"\n                                content_blocks.append(\n                                    ContentBlock(\n                                        tool_call=ToolCall(\n                                            tool=matching_tool,\n                                            tool_call_id=_lstr(\n                                                tool_call.id, origin_trace= origin_id),\n                                            params=json.loads(tool_call.function.arguments),\n                                        )\n                                    )\n                                )\n                                if logger: logger(repr(tool_call))\n                    messages.append(Message(role=role, content=content_blocks))\n            return messages, metadata\n\n\n    # xx: singleton needed\n    openai_provider = OpenAIProvider()\n    register_provider(openai_provider, openai.Client)\nexcept ImportError:\n    pass",
        "filepath": "src\\ell\\providers\\openai.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\openai.py",
            "file_name": "openai.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 1475,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 170,
            "community": null
        },
        "node_id": "providers\\openai.py::1"
    },
    {
        "id": "providers\\openai.py::2",
        "input_type": "file",
        "content": "def _content_block_to_openai_format(content_block: ContentBlock) -> Dict[str, Any]:\n    if (image := content_block.image):\n        image_url = dict(url=serialize_image(image.image) if image.image else image.url)\n        # XXX: Solve per content params better\n        if image.detail: image_url[\"detail\"] = image.detail\n        return {\n            \"type\": \"image_url\",\n            \"image_url\": image_url\n        }\n    elif ((text := content_block.text) is not None): return dict(type=\"text\", text=text)\n    elif (parsed := content_block.parsed): return dict(type=\"text\", text=parsed.model_dump_json())\n    else:\n        raise ValueError(f\"Unsupported content block type for openai: {content_block}\")",
        "filepath": "src\\ell\\providers\\openai.py",
        "metadata": {
            "file_path": "src\\ell\\providers\\openai.py",
            "file_name": "openai.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 158,
            "span_ids": [
                "_content_block_to_openai_format"
            ],
            "start_line": 172,
            "end_line": 184,
            "community": null
        },
        "node_id": "providers\\openai.py::2"
    },
    {
        "id": "ell\\store.py::1",
        "input_type": "file",
        "content": "from abc import ABC, abstractmethod\nfrom contextlib import contextmanager\nfrom datetime import datetime\nfrom typing import Any, Optional, Dict, List, Set, Union\nfrom ell.types._lstr import _lstr\nfrom ell.types import SerializedLMP, Invocation\nfrom ell.types.message import InvocableLM\n\nclass BlobStore(ABC):\n    @abstractmethod\n    def store_blob(self, blob: bytes, blob_id  : str) -> str:\n        \"\"\"Store a blob and return its identifier.\"\"\"\n        pass\n\n    @abstractmethod\n    def retrieve_blob(self, blob_id: str) -> bytes:\n        \"\"\"Retrieve a blob by its identifier.\"\"\"\n        pass",
        "filepath": "src\\ell\\store.py",
        "metadata": {
            "file_path": "src\\ell\\store.py",
            "file_name": "store.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 137,
            "span_ids": [
                "BlobStore.retrieve_blob",
                "BlobStore.store_blob",
                "imports",
                "BlobStore"
            ],
            "start_line": 1,
            "end_line": 18,
            "community": null
        },
        "node_id": "ell\\store.py::1"
    },
    {
        "id": "ell\\store.py::2",
        "input_type": "file",
        "content": "class Store(ABC):\n    \"\"\"\n    Abstract base class for serializers. Defines the interface for serializing and deserializing LMPs and invocations.\n    \"\"\"\n\n    def __init__(self, blob_store: Optional[BlobStore] = None):\n        self.blob_store = blob_store\n\n    @property\n    def has_blob_storage(self) -> bool:\n        return self.blob_store is not None\n\n    @abstractmethod\n    def write_lmp(self, serialized_lmp: SerializedLMP, uses: Dict[str, Any]) -> Optional[Any]:\n        \"\"\"\n        Write an LMP (Language Model Package) to the storage.\n\n        :param serialized_lmp: SerializedLMP object containing all LMP details.\n        :param uses: Dictionary of LMPs used by this LMP.\n        :return: Optional return value.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def write_invocation(self, invocation: Invocation,  consumes: Set[str]) -> Optional[Any]:\n        \"\"\"\n        Write an invocation of an LMP to the storage.\n\n        :param invocation: Invocation object containing all invocation details.\n        :param results: List of SerializedLStr objects representing the results.\n        :param consumes: Set of invocation IDs consumed by this invocation.\n        :return: Optional return value.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def get_cached_invocations(self, lmp_id :str, state_cache_key :str) -> List[Invocation]:\n        \"\"\"\n        Get cached invocations for a given LMP and state cache key.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def get_versions_by_fqn(self, fqn :str) -> List[SerializedLMP]:\n        \"\"\"\n        Get all versions of an LMP by its fully qualified name.\n        \"\"\"\n        pass",
        "filepath": "src\\ell\\store.py",
        "metadata": {
            "file_path": "src\\ell\\store.py",
            "file_name": "store.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 363,
            "span_ids": [
                "Store.write_invocation",
                "Store.has_blob_storage",
                "Store.__init__",
                "Store.get_versions_by_fqn",
                "Store",
                "Store.get_cached_invocations",
                "Store.write_lmp"
            ],
            "start_line": 20,
            "end_line": 67,
            "community": null
        },
        "node_id": "ell\\store.py::2"
    },
    {
        "id": "ell\\store.py::3",
        "input_type": "file",
        "content": "class Store(ABC):\n\n\n    @contextmanager\n    def freeze(self, *lmps: InvocableLM):\n        \"\"\"\n        A context manager for caching operations using a particular store.\n\n        Args:\n            *lmps: InvocableLM objects to freeze.\n\n        Yields:\n            None\n        \"\"\"\n        old_cache_values = {}\n        try:\n            for lmp in lmps:\n                old_cache_values[lmp] = getattr(lmp, '__ell_use_cache__', None)\n                setattr(lmp, '__ell_use_cache__', self)\n            yield\n        finally:\n            # TODO: Implement cache storage logic here\n            for lmp in lmps:\n                if lmp in old_cache_values:\n                    setattr(lmp, '__ell_use_cache__', old_cache_values[lmp])\n                else:\n                    delattr(lmp, '__ell_use_cache__')",
        "filepath": "src\\ell\\store.py",
        "metadata": {
            "file_path": "src\\ell\\store.py",
            "file_name": "store.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 176,
            "span_ids": [
                "Store.freeze"
            ],
            "start_line": 70,
            "end_line": 93,
            "community": null
        },
        "node_id": "ell\\store.py::3"
    },
    {
        "id": "stores\\sql.py::1",
        "input_type": "file",
        "content": "from datetime import datetime, timedelta\nimport json\nimport os\nfrom typing import Any, Optional, Dict, List, Set, Union\nfrom pydantic import BaseModel\nfrom sqlmodel import Session, SQLModel, create_engine, select\nimport ell.store\nimport cattrs\nimport numpy as np\nfrom sqlalchemy.sql import text\nfrom ell.types import InvocationTrace, SerializedLMP, Invocation, InvocationContents\nfrom ell.types._lstr import _lstr\nfrom sqlalchemy import or_, func, and_, extract, FromClause\nfrom sqlalchemy.types import TypeDecorator, VARCHAR\nfrom ell.types.studio import SerializedLMPUses, utc_now\nfrom ell.util.serialization import pydantic_ltype_aware_cattr\nimport gzip\nimport json",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 152,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 18,
            "community": null
        },
        "node_id": "stores\\sql.py::1"
    },
    {
        "id": "stores\\sql.py::2",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n    def __init__(self, db_uri: str, blob_store: Optional[ell.store.BlobStore] = None):\n        self.engine = create_engine(db_uri,\n                                    json_serializer=lambda obj: json.dumps(pydantic_ltype_aware_cattr.unstructure(obj), \n                                     sort_keys=True, default=repr, ensure_ascii=False))\n\n        SQLModel.metadata.create_all(self.engine)\n        self.open_files: Dict[str, Dict[str, Any]] = {}\n        super().__init__(blob_store)\n\n    def write_lmp(self, serialized_lmp: SerializedLMP, uses: Dict[str, Any]) -> Optional[Any]:\n        with Session(self.engine) as session:\n            # Bind the serialized_lmp to the session\n            lmp = session.exec(select(SerializedLMP).filter(SerializedLMP.lmp_id == serialized_lmp.lmp_id)).first()\n\n            if lmp:\n                # Already added to the DB.\n                return lmp\n            else:\n                session.add(serialized_lmp)\n\n            for use_id in uses:\n                used_lmp = session.exec(select(SerializedLMP).where(SerializedLMP.lmp_id == use_id)).first()\n                if used_lmp:\n                    serialized_lmp.uses.append(used_lmp)\n\n            session.commit()\n        return None",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 274,
            "span_ids": [
                "SQLStore",
                "SQLStore.__init__",
                "SQLStore.write_lmp"
            ],
            "start_line": 20,
            "end_line": 47,
            "community": null
        },
        "node_id": "stores\\sql.py::2"
    },
    {
        "id": "stores\\sql.py::3",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n\n    def write_invocation(self, invocation: Invocation, consumes: Set[str]) -> Optional[Any]:\n        with Session(self.engine) as session:\n            lmp = session.exec(select(SerializedLMP).filter(SerializedLMP.lmp_id == invocation.lmp_id)).first()\n            assert lmp is not None, f\"LMP with id {invocation.lmp_id} not found. Writing invocation erroneously\"\n\n            # Increment num_invocations\n            if lmp.num_invocations is None:\n                lmp.num_invocations = 1\n            else:\n                lmp.num_invocations += 1\n\n            # Add the invocation contents\n            session.add(invocation.contents)\n\n            # Add the invocation\n            session.add(invocation)\n\n            # Now create traces.\n            for consumed_id in consumes:\n                session.add(InvocationTrace(\n                    invocation_consumer_id=invocation.id,\n                    invocation_consuming_id=consumed_id\n                ))\n\n            session.commit()\n            return None",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 207,
            "span_ids": [
                "SQLStore.write_invocation"
            ],
            "start_line": 49,
            "end_line": 74,
            "community": null
        },
        "node_id": "stores\\sql.py::3"
    },
    {
        "id": "stores\\sql.py::4",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n\n    def get_cached_invocations(self, lmp_id :str, state_cache_key :str) -> List[Invocation]:\n        with Session(self.engine) as session:\n            return self.get_invocations(session, lmp_filters={\"lmp_id\": lmp_id}, filters={\"state_cache_key\": state_cache_key})\n\n    def get_versions_by_fqn(self, fqn :str) -> List[SerializedLMP]:\n        with Session(self.engine) as session:\n            return self.get_lmps(session, name=fqn)\n\n    ## HELPER METHODS FOR ELL STUDIO! :) \r",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 128,
            "span_ids": [
                "SQLStore.get_cached_invocations",
                "SQLStore.get_versions_by_fqn"
            ],
            "start_line": 76,
            "end_line": 84,
            "community": null
        },
        "node_id": "stores\\sql.py::4"
    },
    {
        "id": "stores\\sql.py::5",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n    def get_latest_lmps(self, session: Session, skip: int = 0, limit: int = 10) -> List[Dict[str, Any]]:\n        \"\"\"\n        Gets all the lmps grouped by unique name with the highest created at\n        \"\"\"\n        subquery = (\n            select(SerializedLMP.name, func.max(SerializedLMP.created_at).label(\"max_created_at\"))\n            .group_by(SerializedLMP.name)\n            .subquery()\n        )\n\n        filters = {\n            \"name\": subquery.c.name,\n            \"created_at\": subquery.c.max_created_at\n        }\n\n        return self.get_lmps(session, skip=skip, limit=limit, subquery=subquery, **filters)",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 158,
            "span_ids": [
                "SQLStore.get_latest_lmps"
            ],
            "start_line": 85,
            "end_line": 100,
            "community": null
        },
        "node_id": "stores\\sql.py::5"
    },
    {
        "id": "stores\\sql.py::6",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n\n\n    def get_lmps(self, session: Session, skip: int = 0, limit: int = 10, subquery=None, **filters: Optional[Dict[str, Any]]) -> List[Dict[str, Any]]:\n\n        query = select(SerializedLMP)\n\n        if subquery is not None:\n            query = query.join(subquery, and_(\n                SerializedLMP.name == subquery.c.name,\n                SerializedLMP.created_at == subquery.c.max_created_at\n            ))\n\n        if filters:\n            for key, value in filters.items():\n                query = query.where(getattr(SerializedLMP, key) == value)\n\n        query = query.order_by(SerializedLMP.created_at.desc())  # Sort by created_at in descending order\n        query = query.offset(skip).limit(limit)\n        results = session.exec(query).all()\n\n        return results",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 188,
            "span_ids": [
                "SQLStore.get_lmps"
            ],
            "start_line": 103,
            "end_line": 121,
            "community": null
        },
        "node_id": "stores\\sql.py::6"
    },
    {
        "id": "stores\\sql.py::7",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n\n    def get_invocations(self, session: Session, lmp_filters: Dict[str, Any], skip: int = 0, limit: int = 10, filters: Optional[Dict[str, Any]] = None, hierarchical: bool = False) -> List[Dict[str, Any]]:\n\n        query = select(Invocation).join(SerializedLMP)\n\n        # Apply LMP filters\n        for key, value in lmp_filters.items():\n            query = query.where(getattr(SerializedLMP, key) == value)\n\n        # Apply invocation filters\n        if filters:\n            for key, value in filters.items():\n                query = query.where(getattr(Invocation, key) == value)\n\n        # Sort from newest to oldest\n        query = query.order_by(Invocation.created_at.desc()).offset(skip).limit(limit)\n\n        invocations = session.exec(query).all()\n        return invocations",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 191,
            "span_ids": [
                "SQLStore.get_invocations"
            ],
            "start_line": 123,
            "end_line": 140,
            "community": null
        },
        "node_id": "stores\\sql.py::7"
    },
    {
        "id": "stores\\sql.py::8",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n\n\n    def get_traces(self, session: Session):\n        query = text(\"\"\"\n        SELECT \n            consumer.lmp_id, \n            trace.*, \n            consumed.lmp_id\n        FROM \n            invocation AS consumer\n        JOIN \n            invocationtrace AS trace ON consumer.id = trace.invocation_consumer_id\n        JOIN \n            invocation AS consumed ON trace.invocation_consuming_id = consumed.id\n        \"\"\")\n        results = session.exec(query).all()\n\n        traces = []\n        for (consumer_lmp_id, consumer_invocation_id, consumed_invocation_id, consumed_lmp_id) in results:\n            traces.append({\n                'consumer': consumer_lmp_id,\n                'consumed': consumed_lmp_id\n            })\n\n        return traces",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 160,
            "span_ids": [
                "SQLStore.get_traces"
            ],
            "start_line": 143,
            "end_line": 165,
            "community": null
        },
        "node_id": "stores\\sql.py::8"
    },
    {
        "id": "stores\\sql.py::9",
        "input_type": "file",
        "content": "class SQLStore(ell.store.Store):\n\n    def get_invocations_aggregate(self, session: Session, lmp_filters: Dict[str, Any] = None, filters: Dict[str, Any] = None, days: int = 30) -> Dict[str, Any]:\n        # Calculate the start date for the graph data\n        start_date = datetime.utcnow() - timedelta(days=days)\n\n        # Base subquery\n        base_subquery = (\n            select(Invocation.created_at, Invocation.latency_ms, Invocation.prompt_tokens, Invocation.completion_tokens, Invocation.lmp_id)\n            .join(SerializedLMP, Invocation.lmp_id == SerializedLMP.lmp_id)\n            .filter(Invocation.created_at >= start_date)\n        )\n\n        # Apply filters\n        if lmp_filters:\n            base_subquery = base_subquery.filter(and_(*[getattr(SerializedLMP, k) == v for k, v in lmp_filters.items()]))\n        if filters:\n            base_subquery = base_subquery.filter(and_(*[getattr(Invocation, k) == v for k, v in filters.items()]))\n\n\n        data = session.exec(base_subquery).all()\n\n        # Calculate aggregate metrics\n        total_invocations = len(data)\n        total_tokens = sum(row.prompt_tokens + row.completion_tokens for row in data)\n        avg_latency = sum(row.latency_ms for row in data) / total_invocations if total_invocations > 0 else 0\n        unique_lmps = len(set(row.lmp_id for row in data))\n\n        # Prepare graph data\n        graph_data = []\n        for row in data:\n            graph_data.append({\n                \"date\": row.created_at,\n                \"avg_latency\": row.latency_ms,\n                \"tokens\": row.prompt_tokens + row.completion_tokens,\n                \"count\": 1\n            })\n\n        return {\n            \"total_invocations\": total_invocations,\n            \"total_tokens\": total_tokens,\n            \"avg_latency\": avg_latency,\n            \"unique_lmps\": unique_lmps,\n            \"graph_data\": graph_data\n        }",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 430,
            "span_ids": [
                "SQLStore.get_invocations_aggregate"
            ],
            "start_line": 167,
            "end_line": 209,
            "community": null
        },
        "node_id": "stores\\sql.py::9"
    },
    {
        "id": "stores\\sql.py::10",
        "input_type": "file",
        "content": "class SQLiteStore(SQLStore):\n    def __init__(self, db_dir: str):\n        assert not db_dir.endswith('.db'), \"Create store with a directory not a db.\"\n\n        os.makedirs(db_dir, exist_ok=True)\n        self.db_dir = db_dir\n        db_path = os.path.join(db_dir, 'ell.db')\n        blob_store = SQLBlobStore(db_dir)\n        super().__init__(f'sqlite:///{db_path}', blob_store=blob_store)\n\nclass SQLBlobStore(ell.store.BlobStore):\n    def __init__(self, db_dir: str):\n        self.db_dir = db_dir\n\n    def store_blob(self, blob: bytes, blob_id  : str) -> str:\n        file_path = self._get_blob_path(blob_id)\n        os.makedirs(os.path.dirname(file_path), exist_ok=True)\n        with gzip.open(file_path, \"wb\") as f:\n            f.write(blob)\n        return blob_id\n\n    def retrieve_blob(self, blob_id: str) -> bytes:\n        file_path = self._get_blob_path(blob_id)\n        with gzip.open(file_path, \"rb\") as f:\n            return f.read()",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 241,
            "span_ids": [
                "SQLBlobStore",
                "SQLBlobStore.__init__",
                "SQLiteStore",
                "SQLBlobStore.store_blob",
                "SQLiteStore.__init__",
                "SQLBlobStore.retrieve_blob"
            ],
            "start_line": 211,
            "end_line": 235,
            "community": null
        },
        "node_id": "stores\\sql.py::10"
    },
    {
        "id": "stores\\sql.py::11",
        "input_type": "file",
        "content": "class SQLBlobStore(ell.store.BlobStore):\n\n    def _get_blob_path(self, id: str, depth: int = 2) -> str:\n        assert \"-\" in id, \"Blob id must have a single - in it to split on.\"\n        _type, _id = id.split(\"-\")\n        increment = 2\n        dirs = [_type] + [_id[i:i+increment] for i in range(0, depth*increment, increment)]\n        file_name = _id[depth*increment:]\n        return os.path.join(self.db_dir, *dirs, file_name)\n\nclass PostgresStore(SQLStore):\n    def __init__(self, db_uri: str):\n        super().__init__(db_uri)",
        "filepath": "src\\ell\\stores\\sql.py",
        "metadata": {
            "file_path": "src\\ell\\stores\\sql.py",
            "file_name": "sql.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 150,
            "span_ids": [
                "PostgresStore.__init__",
                "PostgresStore",
                "SQLBlobStore._get_blob_path"
            ],
            "start_line": 237,
            "end_line": 249,
            "community": null
        },
        "node_id": "stores\\sql.py::11"
    },
    {
        "id": "studio\\__main__.py::1",
        "input_type": "file",
        "content": "import asyncio\nimport logging\nimport socket\nimport time\nimport webbrowser\nimport uvicorn\nfrom argparse import ArgumentParser\nfrom contextlib import closing\nfrom ell.studio.config import Config\nfrom ell.studio.server import create_app\nfrom fastapi.staticfiles import StaticFiles\nfrom fastapi.responses import FileResponse\nfrom pathlib import Path\nfrom watchfiles import awatch\n\n\nlogger = logging.getLogger(__file__)\n\n\ndef _socket_is_open(host, port) -> bool:\n    with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as sock:\n        return sock.connect_ex((host, port)) == 0\n\n\ndef _setup_logging(level):\n    logging.basicConfig(\n        format='%(asctime)s %(levelname)-8s] %(message)s',\n        level=level,\n        datefmt='%Y-%m-%d %H:%M:%S'\n    )",
        "filepath": "src\\ell\\studio\\__main__.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\__main__.py",
            "file_name": "__main__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 176,
            "span_ids": [
                "_socket_is_open",
                "_setup_logging",
                "imports"
            ],
            "start_line": 1,
            "end_line": 30,
            "community": null
        },
        "node_id": "studio\\__main__.py::1"
    },
    {
        "id": "studio\\__main__.py::2",
        "input_type": "file",
        "content": "def main():\n    parser = ArgumentParser(description=\"ell studio\")\n    parser.add_argument(\"--storage-dir\" , default=None,\n                        help=\"Directory for filesystem serializer storage (default: current directory)\")\n    parser.add_argument(\"--pg-connection-string\", default=None,\n                        help=\"PostgreSQL connection string (default: None)\")\n    parser.add_argument(\"--host\", default=\"127.0.0.1\", help=\"Host to run the server on (default: localhost)\")\n    parser.add_argument(\"--port\", type=int, default=5555, help=\"Port to run the server on (default: 5555)\")\n    parser.add_argument(\"--dev\", action=\"store_true\", help=\"Run in development mode\")\n    parser.add_argument(\"--open\", action=\"store_true\", help=\"Opens the studio web UI in a browser\")\n    parser.add_argument(\"--verbose\", \"-v\", action=\"store_true\", help=\"Enables debug logging for more verbose output\")\n    args = parser.parse_args()\n\n    _setup_logging(logging.DEBUG if args.verbose else logging.INFO)\n\n    if args.dev:\n        assert args.port == 5555, \"Port must be 5000 in development mode\"\n\n    config = Config.create(storage_dir=args.storage_dir,\n                    pg_connection_string=args.pg_connection_string)\n    app = create_app(config)\n\n    if not args.dev:\n        # In production mode, serve the built React app\n        static_dir = Path(__file__).parent / \"static\"\n        # app.mount(\"/\", StaticFiles(directory=static_dir, html=True), name=\"static\")\n\n        @app.get(\"/{full_path:path}\")\n        async def serve_react_app(full_path: str):\n            file_path = static_dir / full_path\n            if file_path.exists() and file_path.is_file():\n                return FileResponse(file_path)\n            else:\n                return FileResponse(static_dir / \"index.html\")\n\n    # Respect Config.create behavior, which has fallback to env vars.\n    db_path = Path(config.storage_dir) if config.storage_dir else None\n    # ... other code",
        "filepath": "src\\ell\\studio\\__main__.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\__main__.py",
            "file_name": "__main__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 414,
            "span_ids": [
                "main"
            ],
            "start_line": 33,
            "end_line": 69,
            "community": null
        },
        "node_id": "studio\\__main__.py::2"
    },
    {
        "id": "studio\\__main__.py::3",
        "input_type": "file",
        "content": "def main():\n    # ... other code\n\n    async def db_watcher(db_path, app):\n        last_stat = None\n\n        while True:\n            await asyncio.sleep(0.1)  # Fixed interval of 0.1 seconds\n            try:\n                current_stat = db_path.stat()\n\n                if last_stat is None:\n                    logger.info(f\"Database file found: {db_path}\")\n                    await app.notify_clients(\"database_updated\")\n                else:\n                    # Use a threshold for time comparison to account for filesystem differences\n                    time_threshold = 0.1  # 1 second threshold\n                    time_changed = abs(current_stat.st_mtime - last_stat.st_mtime) > time_threshold\n                    size_changed = current_stat.st_size != last_stat.st_size\n                    inode_changed = current_stat.st_ino != last_stat.st_ino\n\n                    if time_changed or size_changed or inode_changed:\n                        logger.info(\n                            f\"Database changed: mtime {time.ctime(last_stat.st_mtime)} -> {time.ctime(current_stat.st_mtime)}, \"\n                            f\"size {last_stat.st_size} -> {current_stat.st_size}, \"\n                            f\"inode {last_stat.st_ino} -> {current_stat.st_ino}\"\n                        )\n                        await app.notify_clients(\"database_updated\")\n\n                last_stat = current_stat\n            except FileNotFoundError:\n                if last_stat is not None:\n                    logger.info(f\"Database file deleted: {db_path}\")\n                    await app.notify_clients(\"database_updated\")\n                last_stat = None\n                await asyncio.sleep(1)  # Wait a bit longer if the file is missing\n            except Exception as e:\n                logger.info(f\"Error checking database file: {e}\")\n                await asyncio.sleep(1)  # Wait a bit longer on errors\n    # ... other code",
        "filepath": "src\\ell\\studio\\__main__.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\__main__.py",
            "file_name": "__main__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 369,
            "span_ids": [
                "main"
            ],
            "start_line": 71,
            "end_line": 106,
            "community": null
        },
        "node_id": "studio\\__main__.py::3"
    },
    {
        "id": "studio\\__main__.py::4",
        "input_type": "file",
        "content": "def main():\n    # ... other code\n\n    async def open_browser(host, port):\n        while True:\n            logger.debug(f\"Checking TCP port {port} on {host} for readiness.\")\n            if _socket_is_open(host, port):\n                url = f\"http://{host}:{port}\"\n                logger.debug(f\"Port is open, launching {url}.\")\n                webbrowser.open_new(url)\n                return\n\n            logger.debug(f\"Port {port} was not open, retrying.\")\n            await asyncio.sleep(.1)\n\n    # Start the database watcher\n    loop = asyncio.new_event_loop()\n\n    config = uvicorn.Config(app=app, host=args.host, port=args.port, loop=loop)\n    server = uvicorn.Server(config)\n    loop.create_task(server.serve())\n    if db_path:\n        loop.create_task(db_watcher(db_path, app))\n    if args.open:\n        loop.create_task(open_browser(args.host, args.port))\n    loop.run_forever()\n\nif __name__ == \"__main__\":\n    main()",
        "filepath": "src\\ell\\studio\\__main__.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\__main__.py",
            "file_name": "__main__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 211,
            "span_ids": [
                "main",
                "impl:3"
            ],
            "start_line": 108,
            "end_line": 134,
            "community": null
        },
        "node_id": "studio\\__main__.py::4"
    },
    {
        "id": "studio\\config.py::1",
        "input_type": "file",
        "content": "from functools import lru_cache\nimport os\nfrom typing import Optional\nfrom pydantic import BaseModel\n\nimport logging\n\nlogger = logging.getLogger(__name__)\n\n\n# todo. maybe we default storage dir and other things in the future to a well-known location\n# like ~/.ell or something\n@lru_cache\ndef ell_home() -> str:\n    return os.path.join(os.path.expanduser(\"~\"), \".ell\")\n\n\nclass Config(BaseModel):\n    pg_connection_string: Optional[str] = None\n    storage_dir: Optional[str] = None\n\n    @classmethod\n    def create(\n        cls,\n        storage_dir: Optional[str] = None,\n        pg_connection_string: Optional[str] = None,\n    ) -> 'Config':\n        pg_connection_string = pg_connection_string or os.getenv(\"ELL_PG_CONNECTION_STRING\")\n        storage_dir = storage_dir or os.getenv(\"ELL_STORAGE_DIR\")\n\n        # Enforce that we use either sqlite or postgres, but not both\n        if pg_connection_string is not None and storage_dir is not None:\n            raise ValueError(\"Cannot use both sqlite and postgres\")\n\n        # For now, fall back to sqlite if no PostgreSQL connection string is provided\n        if pg_connection_string is None and storage_dir is None:\n            # This intends to honor the default we had set in the CLI\n            storage_dir = os.getcwd()\n\n        return cls(pg_connection_string=pg_connection_string, storage_dir=storage_dir)",
        "filepath": "src\\ell\\studio\\config.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\config.py",
            "file_name": "config.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 291,
            "span_ids": [
                "Config.create",
                "Config",
                "ell_home",
                "imports"
            ],
            "start_line": 1,
            "end_line": 40,
            "community": null
        },
        "node_id": "studio\\config.py::1"
    },
    {
        "id": "studio\\connection_manager.py::1",
        "input_type": "file",
        "content": "from fastapi import WebSocket\n\n\nclass ConnectionManager:\n    def __init__(self):\n        self.active_connections = []\n\n    async def connect(self, websocket: WebSocket):\n        await websocket.accept()\n        self.active_connections.append(websocket)\n\n    def disconnect(self, websocket: WebSocket):\n        self.active_connections.remove(websocket)\n\n    async def broadcast(self, message: str):\n        for connection in self.active_connections:\n            print(f\"Broadcasting message to {connection} {message}\")\n            await connection.send_text(message)",
        "filepath": "src\\ell\\studio\\connection_manager.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\connection_manager.py",
            "file_name": "connection_manager.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 102,
            "span_ids": [
                "ConnectionManager.__init__",
                "ConnectionManager.disconnect",
                "ConnectionManager.broadcast",
                "ConnectionManager.connect",
                "imports",
                "ConnectionManager"
            ],
            "start_line": 1,
            "end_line": 18,
            "community": null
        },
        "node_id": "studio\\connection_manager.py::1"
    },
    {
        "id": "studio\\datamodels.py::1",
        "input_type": "file",
        "content": "from datetime import datetime\nfrom typing import List, Optional, Dict, Any\nfrom sqlmodel import SQLModel\nfrom ell.types import SerializedLMPBase, InvocationBase, InvocationContentsBase\n\n\nclass SerializedLMPWithUses(SerializedLMPBase):\n    lmp_id : str\n    uses: List[SerializedLMPBase]\n\n\nclass InvocationPublic(InvocationBase):\n    lmp: SerializedLMPBase\n    uses: List[\"InvocationPublicWithConsumes\"]\n    contents: InvocationContentsBase\n\nclass InvocationPublicWithConsumes(InvocationPublic):\n    consumes: List[InvocationPublic]\n    consumed_by: List[InvocationPublic]\n\n\n\nfrom pydantic import BaseModel\n\nclass GraphDataPoint(BaseModel):\n    date: datetime\n    count: int\n    avg_latency: float\n    tokens: int\n    # cost: float\n\nclass InvocationsAggregate(BaseModel):\n    total_invocations: int\n    total_tokens: int\n    avg_latency: float\n    # total_cost: float\n    unique_lmps: int\n    # successful_invocations: int\n    # success_rate: float\n    graph_data: List[GraphDataPoint]",
        "filepath": "src\\ell\\studio\\datamodels.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\datamodels.py",
            "file_name": "datamodels.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 235,
            "span_ids": [
                "impl",
                "SerializedLMPWithUses",
                "InvocationsAggregate",
                "InvocationPublic",
                "InvocationPublicWithConsumes",
                "imports",
                "GraphDataPoint"
            ],
            "start_line": 1,
            "end_line": 42,
            "community": null
        },
        "node_id": "studio\\datamodels.py::1"
    },
    {
        "id": "studio\\server.py::1",
        "input_type": "file",
        "content": "from typing import Optional, Dict, Any\n\nfrom sqlmodel import Session\nfrom ell.stores.sql import PostgresStore, SQLiteStore\nfrom ell import __version__\nfrom fastapi import FastAPI, Query, HTTPException, Depends, Response, WebSocket, WebSocketDisconnect\nfrom fastapi.middleware.cors import CORSMiddleware\nimport logging\nimport json\nfrom ell.studio.config import Config\nfrom ell.studio.connection_manager import ConnectionManager\nfrom ell.studio.datamodels import InvocationPublicWithConsumes, SerializedLMPWithUses\n\nfrom ell.types import SerializedLMP\nfrom datetime import datetime, timedelta\nfrom sqlmodel import select\n\n\nlogger = logging.getLogger(__name__)\n\n\nfrom ell.studio.datamodels import InvocationsAggregate\n\n\ndef get_serializer(config: Config):\n    if config.pg_connection_string:\n        return PostgresStore(config.pg_connection_string)\n    elif config.storage_dir:\n        return SQLiteStore(config.storage_dir)\n    else:\n        raise ValueError(\"No storage configuration found\")",
        "filepath": "src\\ell\\studio\\server.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\server.py",
            "file_name": "server.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 200,
            "span_ids": [
                "imports",
                "get_serializer"
            ],
            "start_line": 1,
            "end_line": 31,
            "community": null
        },
        "node_id": "studio\\server.py::1"
    },
    {
        "id": "studio\\server.py::2",
        "input_type": "file",
        "content": "def create_app(config:Config):\n    serializer = get_serializer(config)\n\n    def get_session():\n        with Session(serializer.engine) as session:\n            yield session\n\n    app = FastAPI(title=\"ell Studio\", version=__version__)\n\n    # Enable CORS for all origins\n    app.add_middleware(\n        CORSMiddleware,\n        allow_origins=[\"*\"],\n        allow_credentials=True,\n        allow_methods=[\"*\"],\n        allow_headers=[\"*\"],\n    )\n\n    manager = ConnectionManager()\n\n    @app.websocket(\"/ws\")\n    async def websocket_endpoint(websocket: WebSocket):\n        await manager.connect(websocket)\n        try:\n            while True:\n                data = await websocket.receive_text()\n                # Handle incoming WebSocket messages if needed\n        except WebSocketDisconnect:\n            manager.disconnect(websocket)\n\n\n    @app.get(\"/api/latest/lmps\", response_model=list[SerializedLMPWithUses])\n    def get_latest_lmps(\n        skip: int = Query(0, ge=0),\n        limit: int = Query(100, ge=1, le=100),\n        session: Session = Depends(get_session)\n    ):\n        lmps = serializer.get_latest_lmps(\n            session,\n            skip=skip, limit=limit,\n            )\n        return lmps\n\n    # TOOD: Create a get endpoint to efficient get on the index with /api/lmp/<lmp_id>\n    @app.get(\"/api/lmp/{lmp_id}\")\n    def get_lmp_by_id(lmp_id: str, session: Session = Depends(get_session)):\n        lmp = serializer.get_lmps(session, lmp_id=lmp_id)[0]\n        return lmp\n    # ... other code",
        "filepath": "src\\ell\\studio\\server.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\server.py",
            "file_name": "server.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 341,
            "span_ids": [
                "create_app"
            ],
            "start_line": 35,
            "end_line": 82,
            "community": null
        },
        "node_id": "studio\\server.py::2"
    },
    {
        "id": "studio\\server.py::3",
        "input_type": "file",
        "content": "def create_app(config:Config):\n    # ... other code\n\n\n\n    @app.get(\"/api/lmps\", response_model=list[SerializedLMPWithUses])\n    def get_lmp(\n        lmp_id: Optional[str] = Query(None),\n        name: Optional[str] = Query(None),\n        skip: int = Query(0, ge=0),\n        limit: int = Query(100, ge=1, le=100),\n        session: Session = Depends(get_session)\n    ):\n\n        filters : Dict[str, Any] = {}\n        if name:\n            filters['name'] = name\n        if lmp_id:\n            filters['lmp_id'] = lmp_id\n\n        lmps = serializer.get_lmps(session, skip=skip, limit=limit, **filters)\n\n        if not lmps:\n            raise HTTPException(status_code=404, detail=\"LMP not found\")\n\n        print(lmps[0])\n        return lmps\n\n\n\n    @app.get(\"/api/invocation/{invocation_id}\", response_model=InvocationPublicWithConsumes)\n    def get_invocation(\n        invocation_id: str,\n        session: Session = Depends(get_session)\n    ):\n        invocation = serializer.get_invocations(session, lmp_filters=dict(), filters={\"id\": invocation_id})[0]\n        return invocation\n    # ... other code",
        "filepath": "src\\ell\\studio\\server.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\server.py",
            "file_name": "server.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 275,
            "span_ids": [
                "create_app"
            ],
            "start_line": 86,
            "end_line": 117,
            "community": null
        },
        "node_id": "studio\\server.py::3"
    },
    {
        "id": "studio\\server.py::4",
        "input_type": "file",
        "content": "def create_app(config:Config):\n    # ... other code\n\n    @app.get(\"/api/invocations\", response_model=list[InvocationPublicWithConsumes])\n    def get_invocations(\n        id: Optional[str] = Query(None),\n        hierarchical: Optional[bool] = Query(False),\n        skip: int = Query(0, ge=0),\n        limit: int = Query(100, ge=1, le=100),\n        lmp_name: Optional[str] = Query(None),\n        lmp_id: Optional[str] = Query(None),\n        session: Session = Depends(get_session)\n    ):\n        lmp_filters = {}\n        if lmp_name:\n            lmp_filters[\"name\"] = lmp_name\n        if lmp_id:\n            lmp_filters[\"lmp_id\"] = lmp_id\n\n        invocation_filters = {}\n        if id:\n            invocation_filters[\"id\"] = id\n\n        invocations = serializer.get_invocations(\n            session,\n            lmp_filters=lmp_filters,\n            filters=invocation_filters,\n            skip=skip,\n            limit=limit,\n            hierarchical=hierarchical\n        )\n        return invocations\n    # ... other code",
        "filepath": "src\\ell\\studio\\server.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\server.py",
            "file_name": "server.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 241,
            "span_ids": [
                "create_app"
            ],
            "start_line": 119,
            "end_line": 147,
            "community": null
        },
        "node_id": "studio\\server.py::4"
    },
    {
        "id": "studio\\server.py::5",
        "input_type": "file",
        "content": "def create_app(config:Config):\n    # ... other code\n\n\n    @app.get(\"/api/traces\")\n    def get_consumption_graph(\n        session: Session = Depends(get_session)\n    ):\n        traces = serializer.get_traces(session)\n        return traces\n\n\n\n    @app.get(\"/api/blob/{blob_id}\", response_class=Response)\n    def get_blob(\n        blob_id: str,\n        session: Session = Depends(get_session)\n    ):\n        if serializer.blob_store is None:\n            raise HTTPException(status_code=400, detail=\"Blob storage is not configured\")\n        try:\n            blob_data = serializer.blob_store.retrieve_blob(blob_id)\n            return Response(content=blob_data.decode('utf-8'), media_type=\"application/json\")\n        except FileNotFoundError:\n            raise HTTPException(status_code=404, detail=\"Blob not found\")\n        except Exception as e:\n            logger.error(f\"Error retrieving blob: {str(e)}\")\n            raise HTTPException(status_code=500, detail=\"Internal server error\")\n    # ... other code",
        "filepath": "src\\ell\\studio\\server.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\server.py",
            "file_name": "server.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 209,
            "span_ids": [
                "create_app"
            ],
            "start_line": 150,
            "end_line": 173,
            "community": null
        },
        "node_id": "studio\\server.py::5"
    },
    {
        "id": "studio\\server.py::6",
        "input_type": "file",
        "content": "def create_app(config:Config):\n    # ... other code\n\n    @app.get(\"/api/lmp-history\")\n    def get_lmp_history(\n        days: int = Query(365, ge=1, le=3650),  # Default to 1 year, max 10 years\n        session: Session = Depends(get_session)\n    ):\n        # Calculate the start date\n        start_date = datetime.utcnow() - timedelta(days=days)\n\n        # Query to get all LMP creation times within the date range\n        query = (\n            select(SerializedLMP.created_at)\n            .where(SerializedLMP.created_at >= start_date)\n            .order_by(SerializedLMP.created_at)\n        )\n\n        results = session.exec(query).all()\n\n        # Convert results to a list of dictionaries\n        history = [{\"date\": str(row), \"count\": 1} for row in results]\n\n        return history\n    # ... other code",
        "filepath": "src\\ell\\studio\\server.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\server.py",
            "file_name": "server.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 194,
            "span_ids": [
                "create_app"
            ],
            "start_line": 175,
            "end_line": 195,
            "community": null
        },
        "node_id": "studio\\server.py::6"
    },
    {
        "id": "studio\\server.py::7",
        "input_type": "file",
        "content": "def create_app(config:Config):\n    # ... other code\n\n    async def notify_clients(entity: str, id: Optional[str] = None):\n        message = json.dumps({\"entity\": entity, \"id\": id})\n        await manager.broadcast(message)\n\n    # Add this method to the app object\n    app.notify_clients = notify_clients\n\n\n    @app.get(\"/api/invocations/aggregate\", response_model=InvocationsAggregate)\n    def get_invocations_aggregate(\n        lmp_name: Optional[str] = Query(None),\n        lmp_id: Optional[str] = Query(None),\n        days: int = Query(30, ge=1, le=365),\n        session: Session = Depends(get_session)\n    ):\n        lmp_filters = {}\n        if lmp_name:\n            lmp_filters[\"name\"] = lmp_name\n        if lmp_id:\n            lmp_filters[\"lmp_id\"] = lmp_id\n\n        aggregate_data = serializer.get_invocations_aggregate(session, lmp_filters=lmp_filters, days=days)\n        return InvocationsAggregate(**aggregate_data)\n\n\n\n    return app",
        "filepath": "src\\ell\\studio\\server.py",
        "metadata": {
            "file_path": "src\\ell\\studio\\server.py",
            "file_name": "server.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 227,
            "span_ids": [
                "create_app"
            ],
            "start_line": 197,
            "end_line": 223,
            "community": null
        },
        "node_id": "studio\\server.py::7"
    },
    {
        "id": "types\\__init__.py::1",
        "input_type": "file",
        "content": "\"\"\"\nThe primary types used in ell\n\"\"\"\n\nfrom ell.types.message import *\nfrom ell.types.studio import *\nfrom ell.types._lstr import *",
        "filepath": "src\\ell\\types\\__init__.py",
        "metadata": {
            "file_path": "src\\ell\\types\\__init__.py",
            "file_name": "__init__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 30,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 8,
            "community": null
        },
        "node_id": "types\\__init__.py::1"
    },
    {
        "id": "types\\_lstr.py::1",
        "input_type": "file",
        "content": "\"\"\"\nLM string that supports logits and keeps track of it'sorigin_trace even after mutation.\n\"\"\"\n\nimport numpy as np\nfrom typing import (\n    Optional,\n    Set,\n    SupportsIndex,\n    Union,\n    FrozenSet,\n    Iterable,\n    List,\n    Tuple,\n    Any,\n    Callable,\n)\nfrom typing_extensions import override\nfrom pydantic import BaseModel, GetCoreSchemaHandler\nfrom pydantic_core import CoreSchema\n\nfrom pydantic_core import CoreSchema, core_schema\n\n\nclass _lstr(str):\n    \"\"\"\n     A string class that supports logits and keeps track of itsorigin_trace even after mutation.\n     This class is designed to be used in prompt engineering libraries where it is essential to associate\n     logits with generated text and track the origin of the text.\n\n     The `lstr` class inherits from the built-in `str` class and adds two additional attributes: `logits` and `origin_trace`.\n     The `origin_trace` attribute is a frozen set of strings that represents theorigin_trace(s) of the string.\n\n     The class provides various methods for manipulating the string, such as concatenation, slicing, splitting, and joining.\n     These methods ensure that the logits andorigin_trace(s) are updated correctly based on the operation performed.\n\n     The `lstr` class is particularly useful in LLM libraries for tracing the flow of prompts through various language model calls.\n     By tracking theorigin_trace of each string, it is possible to visualize how outputs from one language model program influence\n     the inputs of another, allowing for a detailed analysis of interactions between different large language models. This capability\n     is crucial for understanding the propagation of prompts in complex LLM workflows and for building visual graphs that depict these interactions.\n\n     It is important to note that any modification to the string (such as concatenation or replacement) will invalidate the associated logits.\n     This is because the logits are specifically tied to the original string content, and any change would require a new computation of logits.\n     The logic behind this is detailed elsewhere in this file.\n\n     Example usage:\n     ```\n     # Create an lstr instance with logits and anorigin_trace\n     logits = np.array([1.0, 2.0, 3.0])\n    origin_trace = \"4e9b7ec9\"\n     lstr_instance = lstr(\"Hello\", logits,origin_trace)\n\n     # Concatenate two lstr instances\n     lstr_instance2 = lstr(\"World\", None, \"7f4d2c3a\")\n     concatenated_lstr = lstr_instance + lstr_instance2\n\n     # Get the logits andorigin_trace of the concatenated lstr\n     print(concatenated_lstr.logits)  # Output: None\n     print(concatenated_lstr.origin_trace)  # Output: frozenset({'4e9b7ec9', '7f4d2c3a'})\n\n     # Split the concatenated lstr into two parts\n     parts = concatenated_lstr.split()\n     print(parts)  # Output: [lstr('Hello', None, frozenset({'4e9b7ec9', '7f4d2c3a'})), lstr('World', None, frozenset({'4e9b7ec9', '7f4d2c3a'}))]\n     ```\n     Attributes:\n        origin_trace (FrozenSet[str]): A frozen set of strings representing theorigin_trace(s) of the string.\n\n     Methods:\n         __new__: Create a new instance of lstr.\n         __repr__: Return a string representation of the lstr instance.\n         __add__: Concatenate this lstr instance with another string or lstr instance.\n         __mod__: Perform a modulo operation between this lstr instance and another string, lstr, or a tuple of strings and lstrs.\n         __mul__: Perform a multiplication operation between this lstr instance and an integer or another lstr.\n         __rmul__: Perform a right multiplication operation between an integer or another lstr and this lstr instance.\n         __getitem__: Get a slice or index of this lstr instance.\n         __getattr__: Get an attribute from this lstr instance.\n         join: Join a sequence of strings or lstr instances into a single lstr instance.\n         split: Split this lstr instance into a list of lstr instances based on a separator.\n         rsplit: Split this lstr instance into a list of lstr instances based on a separator, starting from the right.\n         splitlines: Split this lstr instance into a list of lstr instances based on line breaks.\n         partition: Partition this lstr instance into three lstr instances based on a separator.\n         rpartition: Partition this lstr instance into three lstr instances based on a separator, starting from the right.\n    \"\"\"",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 997,
            "span_ids": [
                "_lstr",
                "docstring"
            ],
            "start_line": 1,
            "end_line": 83,
            "community": null
        },
        "node_id": "types\\_lstr.py::1"
    },
    {
        "id": "types\\_lstr.py::2",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def __new__(\n        cls,\n        content: str,\n        logits: Optional[np.ndarray] = None,\n        origin_trace: Optional[Union[str, FrozenSet[str]]] = None,\n    ):\n        \"\"\"\n         Create a new instance of lstr. The `logits` should be a numpy array and `origin_trace` should be a frozen set of strings or a single string.\n\n         Args:\n         content (str): The string content of the lstr.\n         logits (np.ndarray, optional): The logits associated with this string. Defaults to None.\n        origin_trace (Union[str, FrozenSet[str]], optional): Theorigin_trace(s) of this string. Defaults to None.\n        \"\"\"\n        instance = super(_lstr, cls).__new__(cls, content)\n        # instance._logits = logits\n        if isinstance(origin_trace, str):\n            instance.__origin_trace__ = frozenset({origin_trace})\n        else:\n            instance.__origin_trace__ = (\n                frozenset(origin_trace) if origin_trace is not None else frozenset()\n            )\n        return instance\n\n    # _logits: Optional[np.ndarray]\n    __origin_trace__: FrozenSet[str]",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 251,
            "span_ids": [
                "_lstr.__new__",
                "_lstr:3"
            ],
            "start_line": 85,
            "end_line": 110,
            "community": null
        },
        "node_id": "types\\_lstr.py::2"
    },
    {
        "id": "types\\_lstr.py::3",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @classmethod\n    def __get_pydantic_core_schema__(\n        cls, source_type: Any, handler: GetCoreSchemaHandler\n    ) -> CoreSchema:\n        def validate_lstr(value):\n            if isinstance(value, dict) and value.get(\"__lstr\", False):\n                content = value[\"content\"]\n                origin_trace = value[\"__origin_trace__\"].split(\",\")\n                return cls(content, origin_trace=origin_trace)\n            elif isinstance(value, str):\n                return cls(value)\n            elif isinstance(value, cls):\n                return value\n            else:\n                raise ValueError(f\"Invalid value for lstr: {value}\")\n\n        return core_schema.json_or_python_schema(\n            json_schema=core_schema.typed_dict_schema(\n                {\n                    \"content\": core_schema.typed_dict_field(core_schema.str_schema()),\n                    \"__origin_trace__\": core_schema.typed_dict_field(\n                        core_schema.str_schema()\n                    ),\n                    \"__lstr\": core_schema.typed_dict_field(core_schema.bool_schema()),\n                }\n            ),\n            python_schema=core_schema.union_schema(\n                [\n                    core_schema.is_instance_schema(cls),\n                    core_schema.no_info_plain_validator_function(validate_lstr),\n                ]\n            ),\n            serialization=core_schema.plain_serializer_function_ser_schema(\n                lambda instance: {\n                    \"content\": str(instance),\n                    \"__origin_trace__\": (instance.__origin_trace__),\n                    \"__lstr\": True,\n                }\n            ),\n        )",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 301,
            "span_ids": [
                "_lstr.__get_pydantic_core_schema__"
            ],
            "start_line": 112,
            "end_line": 151,
            "community": null
        },
        "node_id": "types\\_lstr.py::3"
    },
    {
        "id": "types\\_lstr.py::4",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @property\n    def origin_trace(self) -> FrozenSet[str]:\n        \"\"\"\n        Get theorigin_trace(s) of this lstr instance.\n\n        Returns:\n            FrozenSet[str]: A frozen set of strings representing theorigin_trace(s) of this lstr instance.\n        \"\"\"\n        return self.__origin_trace__\n\n    ########################\n    ## Overriding methods ##\n    ########################\n    def __repr__(self) -> str:\n        \"\"\"\n        Return a string representation of this lstr instance.\n\n        Returns:\n            str: A string representation of this lstr instance, including its content, logits, andorigin_trace(s).\n        \"\"\"\n        return super().__repr__()",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 140,
            "span_ids": [
                "_lstr.origin_trace",
                "_lstr.__repr__"
            ],
            "start_line": 153,
            "end_line": 173,
            "community": null
        },
        "node_id": "types\\_lstr.py::4"
    },
    {
        "id": "types\\_lstr.py::5",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def __add__(self, other: Union[str, \"_lstr\"]) -> \"_lstr\":\n        \"\"\"\n        Concatenate this lstr instance with another string or lstr instance.\n\n        Args:\n            other (Union[str, \"lstr\"]): The string or lstr instance to concatenate with this instance.\n\n        Returns:\n            lstr: A new lstr instance containing the concatenated content, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        new_content = super(_lstr, self).__add__(other)\n        self_origin = self.__origin_trace__\n\n        if isinstance(other, _lstr):\n            new_origin = self_origin\n            new_origin = new_origin.union(other.__origin_trace__)\n        else:\n            new_origin = self_origin\n\n        return _lstr(new_content, None, frozenset(new_origin))",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 176,
            "span_ids": [
                "_lstr.__add__"
            ],
            "start_line": 175,
            "end_line": 194,
            "community": null
        },
        "node_id": "types\\_lstr.py::5"
    },
    {
        "id": "types\\_lstr.py::6",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def __mod__(\n        self, other: Union[str, \"_lstr\", Tuple[Union[str, \"_lstr\"], ...]]\n    ) -> \"_lstr\":\n        \"\"\"\n        Perform a modulo operation between this lstr instance and another string, lstr, or a tuple of strings and lstrs,\n        tracing the operation by logging the operands and the result.\n\n        Args:\n            other (Union[str, \"lstr\", Tuple[Union[str, \"lstr\"], ...]]): The right operand in the modulo operation.\n\n        Returns:\n            lstr: A new lstr instance containing the result of the modulo operation, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        # If 'other' is a tuple, we need to handle each element\n        if isinstance(other, tuple):\n            result_content = super(_lstr, self).__mod__(tuple(str(o) for o in other))\n            new__origin_trace__s = set(self.__origin_trace__)\n            for item in other:\n                if isinstance(item, _lstr):\n                    new__origin_trace__s.update(item.__origin_trace__)\n            new__origin_trace__ = frozenset(new__origin_trace__s)\n        else:\n            result_content = super(_lstr, self).__mod__(other)\n            if isinstance(other, _lstr):\n                new__origin_trace__ = self.__origin_trace__.union(\n                    other.__origin_trace__\n                )\n            else:\n                new__origin_trace__ = self.__origin_trace__\n\n        return _lstr(result_content, None, new__origin_trace__)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 331,
            "span_ids": [
                "_lstr.__mod__"
            ],
            "start_line": 196,
            "end_line": 226,
            "community": null
        },
        "node_id": "types\\_lstr.py::6"
    },
    {
        "id": "types\\_lstr.py::7",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def __mul__(self, other: SupportsIndex) -> \"_lstr\":\n        \"\"\"\n        Perform a multiplication operation between this lstr instance and an integer or another lstr,\n        tracing the operation by logging the operands and the result.\n\n        Args:\n            other (Union[SupportsIndex, \"lstr\"]): The right operand in the multiplication operation.\n\n        Returns:\n            lstr: A new lstr instance containing the result of the multiplication operation, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        if isinstance(other, SupportsIndex):\n            result_content = super(_lstr, self).__mul__(other)\n            new__origin_trace__ = self.__origin_trace__\n        else:\n            return NotImplemented\n\n        return _lstr(result_content, None, new__origin_trace__)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 168,
            "span_ids": [
                "_lstr.__mul__"
            ],
            "start_line": 228,
            "end_line": 245,
            "community": null
        },
        "node_id": "types\\_lstr.py::7"
    },
    {
        "id": "types\\_lstr.py::8",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def __rmul__(self, other: SupportsIndex) -> \"_lstr\":\n        \"\"\"\n        Perform a right multiplication operation between an integer or another lstr and this lstr instance,\n        tracing the operation by logging the operands and the result.\n\n        Args:\n            other (Union[SupportsIndex, \"lstr\"]): The left operand in the multiplication operation.\n\n        Returns:\n            lstr: A new lstr instance containing the result of the multiplication operation, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        return self.__mul__(other)  # Multiplication is commutative in this context\r",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 133,
            "span_ids": [
                "_lstr.__rmul__"
            ],
            "start_line": 247,
            "end_line": 258,
            "community": null
        },
        "node_id": "types\\_lstr.py::8"
    },
    {
        "id": "types\\_lstr.py::9",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def __getitem__(self, key: Union[SupportsIndex, slice]) -> \"_lstr\":\n        \"\"\"\n        Get a slice or index of this lstr instance.\n\n        Args:\n            key (Union[SupportsIndex, slice]): The index or slice to retrieve.\n\n        Returns:\n            lstr: A new lstr instance containing the sliced or indexed content, with theorigin_trace(s) preserved.\n        \"\"\"\n        result = super(_lstr, self).__getitem__(key)\n        # This is a matter of opinon. I believe that when you Index into a language model output, you or divorcing the lodges of the indexed result from their contacts which produce them. Therefore, it is only reasonable to directly index into the lodges without changing the original context, and so any mutation on the string should invalidate the logits.\n        # try:\n        #     logit_subset = self._logits[key] if self._logits else None\n        # except:\n        #   logit_subset = None\n        logit_subset = None\n        return _lstr(result, logit_subset, self.__origin_trace__)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 238,
            "span_ids": [
                "_lstr.__getitem__"
            ],
            "start_line": 260,
            "end_line": 277,
            "community": null
        },
        "node_id": "types\\_lstr.py::9"
    },
    {
        "id": "types\\_lstr.py::10",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def __getattribute__(self, name: str) -> Union[Callable, Any]:\n        \"\"\"\n        Get an attribute from this lstr instance.\n\n        Args:\n            name (str): The name of the attribute to retrieve.\n\n        Returns:\n            Union[Callable, Any]: The requested attribute, which may be a method or a value.\n        \"\"\"\n        # Get the attribute from the superclass (str)\n        # First, try to get the attribute from the current class instance\n\n        # Get the attribute using the superclass method\n        attr = super().__getattribute__(name)\n\n        # Check if the attribute is a callable and not defined in lstr class itself\n\n        if name == \"__class__\":\n            return type(self)\n\n        if callable(attr) and name not in _lstr.__dict__:\n            def wrapped(*args: Any, **kwargs: Any) -> Any:\n                result = attr(*args, **kwargs)\n                # If the result is a string, return an lstr instance\n                if isinstance(result, str):\n                    origin_traces = self.__origin_trace__\n                    for arg in args:\n                        if isinstance(arg, _lstr):\n                            origin_traces = origin_traces.union(arg.__origin_trace__)\n                    for key, value in kwargs.items():\n                        if isinstance(value, _lstr):\n                            origin_traces = origin_traces.union(value.__origin_trace__)\n                    return _lstr(result, None, origin_traces)\n\n                return result\n\n            return wrapped\n\n        return attr",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 313,
            "span_ids": [
                "_lstr.__getattribute__"
            ],
            "start_line": 279,
            "end_line": 318,
            "community": null
        },
        "node_id": "types\\_lstr.py::10"
    },
    {
        "id": "types\\_lstr.py::11",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @override\n    def join(self, iterable: Iterable[Union[str, \"_lstr\"]]) -> \"_lstr\":\n        \"\"\"\n        Join a sequence of strings or lstr instances into a single lstr instance.\n\n        Args:\n            iterable (Iterable[Union[str, \"lstr\"]]): The sequence of strings or lstr instances to join.\n\n        Returns:\n            lstr: A new lstr instance containing the joined content, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        new__origin_trace__ = self.__origin_trace__\n        parts = []\n        for item in iterable:\n            if isinstance(item, _lstr):\n                new__origin_trace__ = new__origin_trace__.union(item.__origin_trace__)\n            parts.append(item)\n        new_content = super(_lstr, self).join(parts)\n\n        return _lstr(new_content, None, new__origin_trace__)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 189,
            "span_ids": [
                "_lstr.join"
            ],
            "start_line": 320,
            "end_line": 339,
            "community": null
        },
        "node_id": "types\\_lstr.py::11"
    },
    {
        "id": "types\\_lstr.py::12",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @override\n    def split(\n        self, sep: Optional[Union[str, \"_lstr\"]] = None, maxsplit: SupportsIndex = -1\n    ) -> List[\"_lstr\"]:\n        \"\"\"\n        Split this lstr instance into a list of lstr instances based on a separator.\n\n        Args:\n            sep (Optional[Union[str, \"lstr\"]], optional): The separator to split on. Defaults to None.\n            maxsplit (SupportsIndex, optional): The maximum number of splits to perform. Defaults to -1.\n\n        Returns:\n            List[\"lstr\"]: A list of lstr instances containing the split content, with theorigin_trace(s) preserved.\n        \"\"\"\n        return self._split_helper(super(_lstr, self).split, sep, maxsplit)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 169,
            "span_ids": [
                "_lstr.split"
            ],
            "start_line": 341,
            "end_line": 355,
            "community": null
        },
        "node_id": "types\\_lstr.py::12"
    },
    {
        "id": "types\\_lstr.py::13",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @override\n    def rsplit(\n        self, sep: Optional[Union[str, \"_lstr\"]] = None, maxsplit: SupportsIndex = -1\n    ) -> List[\"_lstr\"]:\n        \"\"\"\n        Split this lstr instance into a list of lstr instances based on a separator, starting from the right.\n\n        Args:\n            sep (Optional[Union[str, \"lstr\"]], optional): The separator to split on. Defaults to None.\n            maxsplit (SupportsIndex, optional): The maximum number of splits to perform. Defaults to -1.\n\n        Returns:\n            List[\"lstr\"]: A list of lstr instances containing the split content, with theorigin_trace(s) preserved.\n        \"\"\"\n        return self._split_helper(super(_lstr, self).rsplit, sep, maxsplit)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 176,
            "span_ids": [
                "_lstr.rsplit"
            ],
            "start_line": 357,
            "end_line": 371,
            "community": null
        },
        "node_id": "types\\_lstr.py::13"
    },
    {
        "id": "types\\_lstr.py::14",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @override\n    def splitlines(self, keepends: bool = False) -> List[\"_lstr\"]:\n        \"\"\"\n        Split this lstr instance into a list of lstr instances based on line breaks.\n\n        Args:\n            keepends (bool, optional): Whether to include the line breaks in the resulting lstr instances. Defaults to False.\n\n        Returns:\n            List[\"lstr\"]: A list of lstr instances containing the split content, with theorigin_trace(s) preserved.\n        \"\"\"\n        return [\n            _lstr(p, None, self.__origin_trace__)\n            for p in super(_lstr, self).splitlines(keepends=keepends)\n        ]",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 146,
            "span_ids": [
                "_lstr.splitlines"
            ],
            "start_line": 373,
            "end_line": 387,
            "community": null
        },
        "node_id": "types\\_lstr.py::14"
    },
    {
        "id": "types\\_lstr.py::15",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @override\n    def partition(self, sep: Union[str, \"_lstr\"]) -> Tuple[\"_lstr\", \"_lstr\", \"_lstr\"]:\n        \"\"\"\n        Partition this lstr instance into three lstr instances based on a separator.\n\n        Args:\n            sep (Union[str, \"lstr\"]): The separator to partition on.\n\n        Returns:\n            Tuple[\"lstr\", \"lstr\", \"lstr\"]: A tuple of three lstr instances containing the content before the separator, the separator itself, and the content after the separator, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        return self._partition_helper(super(_lstr, self).partition, sep)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 146,
            "span_ids": [
                "_lstr.partition"
            ],
            "start_line": 389,
            "end_line": 400,
            "community": null
        },
        "node_id": "types\\_lstr.py::15"
    },
    {
        "id": "types\\_lstr.py::16",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    @override\n    def rpartition(self, sep: Union[str, \"_lstr\"]) -> Tuple[\"_lstr\", \"_lstr\", \"_lstr\"]:\n        \"\"\"\n        Partition this lstr instance into three lstr instances based on a separator, starting from the right.\n\n        Args:\n            sep (Union[str, \"lstr\"]): The separator to partition on.\n\n        Returns:\n            Tuple[\"lstr\", \"lstr\", \"lstr\"]: A tuple of three lstr instances containing the content before the separator, the separator itself, and the content after the separator, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        return self._partition_helper(super(_lstr, self).rpartition, sep)",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 153,
            "span_ids": [
                "_lstr.rpartition"
            ],
            "start_line": 402,
            "end_line": 413,
            "community": null
        },
        "node_id": "types\\_lstr.py::16"
    },
    {
        "id": "types\\_lstr.py::17",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def _partition_helper(\n        self, method, sep: Union[str, \"_lstr\"]\n    ) -> Tuple[\"_lstr\", \"_lstr\", \"_lstr\"]:\n        \"\"\"\n        Helper method for partitioning this lstr instance based on a separator.\n\n        Args:\n            method (Callable): The partitioning method to use (either partition or rpartition).\n            sep (Union[str, \"lstr\"]): The separator to partition on.\n\n        Returns:\n            Tuple[\"lstr\", \"lstr\", \"lstr\"]: A tuple of three lstr instances containing the content before the separator, the separator itself, and the content after the separator, with theorigin_trace(s) updated accordingly.\n        \"\"\"\n        part1, part2, part3 = method(sep)\n        new__origin_trace__ = (\n            self.__origin_trace__ | sep.__origin_trace__\n            if isinstance(sep, _lstr)\n            else self.__origin_trace__\n        )\n        return (\n            _lstr(part1, None, new__origin_trace__),\n            _lstr(part2, None, new__origin_trace__),\n            _lstr(part3, None, new__origin_trace__),\n        )",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 253,
            "span_ids": [
                "_lstr._partition_helper"
            ],
            "start_line": 415,
            "end_line": 438,
            "community": null
        },
        "node_id": "types\\_lstr.py::17"
    },
    {
        "id": "types\\_lstr.py::18",
        "input_type": "file",
        "content": "class _lstr(str):\n\n    def _split_helper(\n        self,\n        method,\n        sep: Optional[Union[str, \"_lstr\"]] = None,\n        maxsplit: SupportsIndex = -1,\n    ) -> List[\"_lstr\"]:\n        \"\"\"\n        Helper method for splitting this lstr instance based on a separator.\n\n        Args:\n            method (Callable): The splitting method to use (either split or rsplit).\n            sep (Optional[Union[str, \"lstr\"]], optional): The separator to split on. Defaults to None.\n            maxsplit (SupportsIndex, optional): The maximum number of splits to perform. Defaults to -1.\n\n        Returns:\n            List[\"lstr\"]: A list of lstr instances containing the split content, with theorigin_trace(s) preserved.\n        \"\"\"\n        origin_traces = (\n            self.__origin_trace__ | sep.__origin_trace__\n            if isinstance(sep, _lstr)\n            else self.__origin_trace__\n        )\n        parts = method(sep, maxsplit)\n        return [_lstr(part, None, origin_traces) for part in parts]",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 230,
            "span_ids": [
                "_lstr._split_helper"
            ],
            "start_line": 440,
            "end_line": 463,
            "community": null
        },
        "node_id": "types\\_lstr.py::18"
    },
    {
        "id": "types\\_lstr.py::19",
        "input_type": "file",
        "content": "if __name__ == \"__main__\":\n    import timeit\n    import random\n    import string\n\n    def generate_random_string(length):\n        return \"\".join(random.choices(string.ascii_letters + string.digits, k=length))\n\n    def test_concatenation():\n        s1 = generate_random_string(1000)\n        s2 = generate_random_string(1000)\n\n        lstr_time = timeit.timeit(lambda: _lstr(s1) + _lstr(s2), number=10000)\n        str_time = timeit.timeit(lambda: s1 + s2, number=10000)\n\n        print(f\"Concatenation: lstr: {lstr_time:.6f}s, str: {str_time:.6f}s\")\n\n    def test_slicing():\n        s = generate_random_string(10000)\n        ls = _lstr(s)\n\n        lstr_time = timeit.timeit(lambda: ls[1000:2000], number=10000)\n        str_time = timeit.timeit(lambda: s[1000:2000], number=10000)\n\n        print(f\"Slicing: lstr: {lstr_time:.6f}s, str: {str_time:.6f}s\")\n\n    def test_splitting():\n        s = generate_random_string(10000)\n        ls = _lstr(s)\n\n        lstr_time = timeit.timeit(lambda: ls.split(), number=1000)\n        str_time = timeit.timeit(lambda: s.split(), number=1000)\n\n        print(f\"Splitting: lstr: {lstr_time:.6f}s, str: {str_time:.6f}s\")\n\n    def test_joining():\n        words = [generate_random_string(10) for _ in range(1000)]\n        lwords = [_lstr(word) for word in words]\n\n        lstr_time = timeit.timeit(lambda: _lstr(\" \").join(lwords), number=1000)\n        str_time = timeit.timeit(lambda: \" \".join(words), number=1000)\n\n        print(f\"Joining: lstr: {lstr_time:.6f}s, str: {str_time:.6f}s\")\n\n    print(\"Running performance tests...\")\n    test_concatenation()\n    test_slicing()\n    test_splitting()\n    test_joining()\n\n    import cProfile\n    import pstats\n    from io import StringIO\n\n    def test_add():\n        s1 = generate_random_string(1000)\n        s2 = generate_random_string(1000)\n        ls1 = _lstr(s1, None, \"origin1\")\n        ls2 = _lstr(s2, None, \"origin2\")\n\n        for _ in range(100000):\n            result = ls1 + ls2\n\n    print(\"\\nProfiling __add__ method:\")\n    profiler = cProfile.Profile()\n    profiler.enable()\n    test_add()\n    profiler.disable()\n\n    s = StringIO()\n    ps = pstats.Stats(profiler, stream=s).sort_stats(\"cumulative\")\n    ps.print_stats(20)  # Print top 20 lines\n    print(s.getvalue())",
        "filepath": "src\\ell\\types\\_lstr.py",
        "metadata": {
            "file_path": "src\\ell\\types\\_lstr.py",
            "file_name": "_lstr.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 650,
            "span_ids": [
                "impl"
            ],
            "start_line": 466,
            "end_line": 539,
            "community": null
        },
        "node_id": "types\\_lstr.py::19"
    },
    {
        "id": "types\\message.py::1",
        "input_type": "file",
        "content": "# todo: implement tracing for structured outs. this a v2 feature.\nimport json\nfrom ell.types._lstr import _lstr\nfrom functools import cached_property\nimport numpy as np\nimport base64\nfrom io import BytesIO\nfrom PIL import Image as PILImage\n\nfrom pydantic import BaseModel, ConfigDict, model_validator, field_serializer\nfrom sqlmodel import Field\n\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\n\nfrom typing import Any, Callable, Dict, List, Optional, Union\n\nfrom ell.util.serialization import serialize_image\n_lstr_generic = Union[_lstr, str]\nInvocableTool = Callable[..., Union[\"ToolResult\", _lstr_generic, List[\"ContentBlock\"], ]]\n\n# AnyContent represents any type that can be passed to Message.\nAnyContent = Union[\"ContentBlock\", str, \"ToolCall\", \"ToolResult\", \"ImageContent\", np.ndarray, PILImage.Image, BaseModel]",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 193,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 22,
            "community": null
        },
        "node_id": "types\\message.py::1"
    },
    {
        "id": "types\\message.py::2",
        "input_type": "file",
        "content": "class ToolResult(BaseModel):\n    tool_call_id: _lstr_generic\n    result: List[\"ContentBlock\"]\n\n    @property\n    def text(self) -> str:\n        return _content_to_text(self.result)\n\n    @property\n    def text_only(self) -> str:\n        return _content_to_text_only(self.result)\n\n    # # XXX: Possibly deprecate\n    # def readable_repr(self) -> str:\n    #     return f\"ToolResult(tool_call_id={self.tool_call_id}, result={_content_to_text(self.result)})\"\n\n    def __repr__(self):\n        return f\"{self.__class__.__name__}(tool_call_id={self.tool_call_id}, result={_content_to_text(self.result)})\"",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 153,
            "span_ids": [
                "ToolResult.__repr__",
                "ToolResult",
                "ToolResult.text",
                "ToolResult.text_only"
            ],
            "start_line": 25,
            "end_line": 42,
            "community": null
        },
        "node_id": "types\\message.py::2"
    },
    {
        "id": "types\\message.py::3",
        "input_type": "file",
        "content": "class ToolCall(BaseModel):\n    tool : InvocableTool\n    tool_call_id : Optional[_lstr_generic] = Field(default=None)\n    params : BaseModel\n\n    def __init__(self, tool, params : Union[BaseModel, Dict[str, Any]],  tool_call_id=None):\n        if not isinstance(params, BaseModel):\n            params = tool.__ell_params_model__(**params) #convenience.\n        super().__init__(tool=tool, tool_call_id=tool_call_id, params=params)\n\n    def __call__(self, **kwargs):\n        assert not kwargs, \"Unexpected arguments provided. Calling a tool uses the params provided in the ToolCall.\"\n\n        # XXX: TODO: MOVE TRACKING CODE TO _TRACK AND OUT OF HERE AND API.\n        return self.tool(**self.params.model_dump())\n\n    # XXX: Deprecate in 0.1.0\n    def call_and_collect_as_message_block(self):\n        raise DeprecationWarning(\"call_and_collect_as_message_block is deprecated. Use collect_as_content_block instead.\")\n\n    def call_and_collect_as_content_block(self):\n        res = self.tool(**self.params.model_dump(), _tool_call_id=self.tool_call_id)\n        return ContentBlock(tool_result=res)\n\n    def call_and_collect_as_message(self):\n        return Message(role=\"user\", content=[self.call_and_collect_as_message_block()])\n\n    def __repr__(self):\n        return f\"{self.__class__.__name__}({self.tool.__name__}({self.params}), tool_call_id='{self.tool_call_id}')\"",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 320,
            "span_ids": [
                "ToolCall.call_and_collect_as_message",
                "ToolCall.call_and_collect_as_content_block",
                "ToolCall",
                "ToolCall.__call__",
                "ToolCall.__repr__",
                "ToolCall.__init__",
                "ToolCall.call_and_collect_as_message_block"
            ],
            "start_line": 44,
            "end_line": 72,
            "community": null
        },
        "node_id": "types\\message.py::3"
    },
    {
        "id": "types\\message.py::4",
        "input_type": "file",
        "content": "class ImageContent(BaseModel):\n    model_config = ConfigDict(arbitrary_types_allowed=True)\n\n    image: Optional[PILImage.Image] = Field(default=None)\n    url: Optional[str] = Field(default=None)\n    detail: Optional[str] = Field(default=None)\n\n    @model_validator(mode='after')\n    def check_image_or_url(self):\n        if self.image is not None and self.url is not None:\n            raise ValueError(\"Both 'image' and 'url' cannot be set simultaneously.\")\n        if self.image is None and self.url is None:\n            raise ValueError(\"Either 'image' or 'url' must be set.\")\n        return self",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 132,
            "span_ids": [
                "ImageContent.check_image_or_url",
                "ImageContent"
            ],
            "start_line": 75,
            "end_line": 88,
            "community": null
        },
        "node_id": "types\\message.py::4"
    },
    {
        "id": "types\\message.py::5",
        "input_type": "file",
        "content": "class ImageContent(BaseModel):\n\n    @classmethod\n    def coerce(cls, value: Union[str, np.ndarray, PILImage.Image, \"ImageContent\"]):\n        if isinstance(value, cls):\n            return value\n\n        if isinstance(value, str):\n            if value.startswith('http://') or value.startswith('https://'):\n                return cls(url=value)\n            try:\n                img_data = base64.b64decode(value)\n                img = PILImage.open(BytesIO(img_data))\n                if img.mode not in ('L', 'RGB', 'RGBA'):\n                    return cls(image=img.convert('RGB'))\n            except:\n                raise ValueError(\"Invalid base64 string or URL for image\")\n\n        if isinstance(value, np.ndarray):\n            if value.ndim == 3 and value.shape[2] in (3, 4):\n                mode = 'RGB' if value.shape[2] == 3 else 'RGBA'\n                return cls(image=PILImage.fromarray(value, mode=mode))\n            else:\n                raise ValueError(f\"Invalid numpy array shape for image: {value.shape}. Expected 3D array with 3 or 4 channels.\")\n\n        if isinstance(value, PILImage.Image):\n            if value.mode not in ('L', 'RGB', 'RGBA'):\n                value = value.convert('RGB')\n            return cls(image=value)\n\n        raise ValueError(f\"Invalid image type: {type(value)}\")\n\n    @field_serializer('image')\n    def serialize_image(self, image: Optional[PILImage.Image], _info):\n        if image is None:\n            return None\n        return serialize_image(image)",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 322,
            "span_ids": [
                "ImageContent.coerce",
                "ImageContent.serialize_image"
            ],
            "start_line": 90,
            "end_line": 124,
            "community": null
        },
        "node_id": "types\\message.py::5"
    },
    {
        "id": "types\\message.py::6",
        "input_type": "file",
        "content": "class ContentBlock(BaseModel):\n    model_config = ConfigDict(arbitrary_types_allowed=True)\n\n    text: Optional[_lstr_generic] = Field(default=None)\n    image: Optional[ImageContent] = Field(default=None)\n    audio: Optional[Union[np.ndarray, List[float]]] = Field(default=None)\n    tool_call: Optional[ToolCall] = Field(default=None)\n    parsed: Optional[BaseModel] = Field(default=None)\n    tool_result: Optional[ToolResult] = Field(default=None)\n    # TODO: Add a JSON type? This would be nice for response_format. This is different than resposne_format = model. Or we could be opinionated and automatically parse the json response. That might be nice.\n    # This breaks us maintaing parity with the openai python client in some sen but so does image.\n\n    def __init__(self, *args, **kwargs):\n        if \"image\" in kwargs and not isinstance(kwargs[\"image\"], ImageContent):\n            im = kwargs[\"image\"] = ImageContent.coerce(kwargs[\"image\"])\n            # XXX: Backwards compatibility, Deprecate.\n            if (d := kwargs.get(\"image_detail\", None)): im.detail = d\n\n        super().__init__(*args, **kwargs)\n\n\n    @model_validator(mode='after')\n    def check_single_non_null(self):\n        non_null_fields = [field for field, value in self.__dict__.items() if value is not None]\n        if len(non_null_fields) > 1:\n            raise ValueError(f\"Only one field can be non-null. Found: {', '.join(non_null_fields)}\")\n        return self\n\n    def __str__(self):\n        return repr(self)\n\n    def __repr__(self):\n        non_null_fields = [f\"{field}={value}\" for field, value in self.__dict__.items() if value is not None]\n        return f\"ContentBlock({', '.join(non_null_fields)})\"\n\n    @property\n    def type(self):\n        if self.text is not None:\n            return \"text\"\n        if self.image is not None:\n            return \"image\"\n        if self.audio is not None:\n            return \"audio\"\n        if self.tool_call is not None:\n            return \"tool_call\"\n        if self.parsed is not None:\n            return \"parsed\"\n        if self.tool_result is not None:\n            return \"tool_result\"\n        return None\n\n    @property\n    def content(self):\n        return getattr(self, self.type)",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 516,
            "span_ids": [
                "ContentBlock.__repr__",
                "ContentBlock.type",
                "ContentBlock.check_single_non_null",
                "ContentBlock.__init__",
                "ContentBlock.__str__",
                "ContentBlock.content",
                "ContentBlock"
            ],
            "start_line": 126,
            "end_line": 179,
            "community": null
        },
        "node_id": "types\\message.py::6"
    },
    {
        "id": "types\\message.py::7",
        "input_type": "file",
        "content": "class ContentBlock(BaseModel):\n\n    @classmethod\n    def coerce(cls, content: AnyContent) -> \"ContentBlock\":\n        \"\"\"\n        Coerce various types of content into a ContentBlock.\n\n        This method provides a flexible way to create ContentBlock instances from different types of input.\n\n        Args:\n        content: The content to be coerced into a ContentBlock. Can be one of the following types:\n        - str: Will be converted to a text ContentBlock.\n        - ToolCall: Will be converted to a tool_call ContentBlock.\n        - ToolResult: Will be converted to a tool_result ContentBlock.\n        - BaseModel: Will be converted to a parsed ContentBlock.\n        - ContentBlock: Will be returned as-is.\n        - Image: Will be converted to an image ContentBlock.\n        - np.ndarray: Will be converted to an image ContentBlock.\n        - PILImage.Image: Will be converted to an image ContentBlock.\n\n        Returns:\n        ContentBlock: A new ContentBlock instance containing the coerced content.\n\n        Raises:\n        ValueError: If the content cannot be coerced into a valid ContentBlock.\n\n        Examples:\n        >>> ContentBlock.coerce(\"Hello, world!\")\n        ContentBlock(text=\"Hello, world!\")\n\n        >>> tool_call = ToolCall(...)\n        >>> ContentBlock.coerce(tool_call)\n        ContentBlock(tool_call=tool_call)\n\n        >>> tool_result = ToolResult(...)\n        >>> ContentBlock.coerce(tool_result)\n        ContentBlock(tool_result=tool_result)\n\n        >>> class MyModel(BaseModel):\n        ...     field: str\n        >>> model_instance = MyModel(field=\"value\")\n        >>> ContentBlock.coerce(model_instance)\n        ContentBlock(parsed=model_instance)\n\n        >>> from PIL import Image as PILImage\n        >>> img = PILImage.new('RGB', (100, 100))\n        >>> ContentBlock.coerce(img)\n        ContentBlock(image=ImageContent(image=<PIL.Image.Image object>))\n\n        >>> import numpy as np\n        >>> arr = np.random.rand(100, 100, 3)\n        >>> ContentBlock.coerce(arr)\n        ContentBlock(image=ImageContent(image=<PIL.Image.Image object>))\n\n        >>> image = Image(url=\"https://example.com/image.jpg\")\n        >>> ContentBlock.coerce(image)\n        ContentBlock(image=ImageContent(url=\"https://example.com/image.jpg\"))\n\n        Notes:\n        - This method is particularly useful when working with heterogeneous content types\n          and you want to ensure they are all properly encapsulated in ContentBlock instances.\n        - The method performs type checking and appropriate conversions to ensure the resulting\n          ContentBlock is valid according to the model's constraints.\n        - For image content, Image objects, PIL Image objects, and numpy arrays are supported,\n          with automatic conversion to the appropriate format.\n        - As a last resort, the method will attempt to create an image from the input before\n          raising a ValueError.\n        \"\"\"\n        if isinstance(content, ContentBlock):\n            return content\n        if isinstance(content, str):\n            return cls(text=content)\n        if isinstance(content, ToolCall):\n            return cls(tool_call=content)\n        if isinstance(content, ToolResult):\n            return cls(tool_result=content)\n        if isinstance(content, (ImageContent, np.ndarray, PILImage.Image)):\n            return cls(image=ImageContent.coerce(content))\n        if isinstance(content, BaseModel):\n            return cls(parsed=content)\n\n        raise ValueError(f\"Invalid content type: {type(content)}\")\n\n    @field_serializer('parsed')\n    def serialize_parsed(self, value: Optional[BaseModel], _info):\n        if value is None:\n            return None\n        return value.model_dump(exclude_none=True, exclude_unset=True)",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 760,
            "span_ids": [
                "ContentBlock.serialize_parsed",
                "ContentBlock.coerce"
            ],
            "start_line": 181,
            "end_line": 266,
            "community": null
        },
        "node_id": "types\\message.py::7"
    },
    {
        "id": "types\\message.py::8",
        "input_type": "file",
        "content": "def to_content_blocks(\n    content: Optional[Union[AnyContent, List[AnyContent]]] = None,\n    **content_block_kwargs\n) -> List[ContentBlock]:\n    \"\"\"\n    Coerce a variety of input types into a list of ContentBlock objects.\n\n    Args:\n    content: The content to be coerced. Can be a single item or a list of items.\n             Supported types include str, ContentBlock, ToolCall, ToolResult, BaseModel, Image, np.ndarray, and PILImage.Image.\n    **content_block_kwargs: Additional keyword arguments to pass to ContentBlock creation if content is None.\n\n    Returns:\n    List[ContentBlock]: A list of ContentBlock objects created from the input content.\n\n    Examples:\n    >>> coerce_content_list(\"Hello\")\n    [ContentBlock(text=\"Hello\")]\n\n    >>> coerce_content_list([ContentBlock(text=\"Hello\"), \"World\"])\n    [ContentBlock(text=\"Hello\"), ContentBlock(text=\"World\")]\n\n    >>> from PIL import Image as PILImage\n    >>> pil_image = PILImage.new('RGB', (100, 100))\n    >>> coerce_content_list(pil_image)\n    [ContentBlock(image=Image(image=<PIL.Image.Image object>))]\n\n    >>> coerce_content_list(Image(url=\"https://example.com/image.jpg\"))\n    [ContentBlock(image=Image(url=\"https://example.com/image.jpg\"))]\n\n    >>> coerce_content_list(None, text=\"Default text\")\n    [ContentBlock(text=\"Default text\")]\n    \"\"\"\n    if content is None:\n        return [ContentBlock(**content_block_kwargs)]\n\n    if not isinstance(content, list):\n        content = [content]\n\n    return [ContentBlock.model_validate(ContentBlock.coerce(c)) for c in content]",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 351,
            "span_ids": [
                "to_content_blocks"
            ],
            "start_line": 269,
            "end_line": 308,
            "community": null
        },
        "node_id": "types\\message.py::8"
    },
    {
        "id": "types\\message.py::9",
        "input_type": "file",
        "content": "class Message(BaseModel):\n    role: str\n    content: List[ContentBlock]\n\n\n    def __init__(self, role: str, content: Union[AnyContent, List[AnyContent], None] = None, **content_block_kwargs):\n        content_blocks = to_content_blocks(content, **content_block_kwargs)\n\n        super().__init__(role=role, content=content_blocks)\n\n    # XXX: This choice of naming is unfortunate, but it is what it is.\n    @property\n    def text(self) -> str:\n        \"\"\"Returns all text content, replacing non-text content with their representations.\n\n        Example:\n            >>> message = Message(role=\"user\", content=[\"Hello\", PILImage.new('RGB', (100, 100)), \"World\"])\n            >>> message.text\n            'Hello\\\\n<PilImage>\\\\nWorld'\n        \"\"\"\n        return _content_to_text(self.content)",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 182,
            "span_ids": [
                "Message",
                "Message.text",
                "Message.__init__"
            ],
            "start_line": 312,
            "end_line": 332,
            "community": null
        },
        "node_id": "types\\message.py::9"
    },
    {
        "id": "types\\message.py::10",
        "input_type": "file",
        "content": "class Message(BaseModel):\n\n    @property\n    def images(self) -> List[ImageContent]:\n        \"\"\"Returns a list of all image content.\n\n        Example:\n            >>> from PIL import Image as PILImage\n            >>> image1 = Image(url=\"https://example.com/image.jpg\")\n            >>> image2 = Image(image=PILImage.new('RGB', (200, 200)))\n            >>> message = Message(role=\"user\", content=[\"Text\", image1, \"More text\", image2])\n            >>> len(message.images)\n            2\n            >>> isinstance(message.images[0], Image)\n            True\n            >>> message.images[0].url\n            'https://example.com/image.jpg'\n            >>> isinstance(message.images[1].image, PILImage.Image)\n            True\n        \"\"\"\n        return [c.image for c in self.content if c.image]",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 175,
            "span_ids": [
                "Message.images"
            ],
            "start_line": 334,
            "end_line": 352,
            "community": null
        },
        "node_id": "types\\message.py::10"
    },
    {
        "id": "types\\message.py::11",
        "input_type": "file",
        "content": "class Message(BaseModel):\n\n    @property\n    def audios(self) -> List[Union[np.ndarray, List[float]]]:\n        \"\"\"Returns a list of all audio content.\n\n        Example:\n            >>> audio1 = np.array([0.1, 0.2, 0.3])\n            >>> audio2 = np.array([0.4, 0.5, 0.6])\n            >>> message = Message(role=\"user\", content=[\"Text\", audio1, \"More text\", audio2])\n            >>> len(message.audios)\n            2\n        \"\"\"\n        return [c.audio for c in self.content if c.audio]",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 133,
            "span_ids": [
                "Message.audios"
            ],
            "start_line": 354,
            "end_line": 365,
            "community": null
        },
        "node_id": "types\\message.py::11"
    },
    {
        "id": "types\\message.py::12",
        "input_type": "file",
        "content": "class Message(BaseModel):\n\n    @property\n    def text_only(self) -> str:\n        \"\"\"Returns only the text content, ignoring non-text content.\n\n        Example:\n            >>> message = Message(role=\"user\", content=[\"Hello\", PILImage.new('RGB', (100, 100)), \"World\"])\n            >>> message.text_only\n            'Hello\\\\nWorld'\n        \"\"\"\n        return _content_to_text_only(self.content)\n\n    @cached_property\n    def tool_calls(self) -> List[ToolCall]:\n        \"\"\"Returns a list of all tool calls.\n\n        Example:\n            >>> tool_call = ToolCall(tool=lambda x: x, params=BaseModel())\n            >>> message = Message(role=\"user\", content=[\"Text\", tool_call])\n            >>> len(message.tool_calls)\n            1\n        \"\"\"\n        return [c.tool_call for c in self.content if c.tool_call is not None]\n\n    @property\n    def tool_results(self) -> List[ToolResult]:\n        \"\"\"Returns a list of all tool results.\n\n        Example:\n            >>> tool_result = ToolResult(tool_call_id=\"123\", result=[ContentBlock(text=\"Result\")])\n            >>> message = Message(role=\"user\", content=[\"Text\", tool_result])\n            >>> len(message.tool_results)\n            1\n        \"\"\"\n        return [c.tool_result for c in self.content if c.tool_result is not None]",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 282,
            "span_ids": [
                "Message.tool_results",
                "Message.text_only",
                "Message.tool_calls"
            ],
            "start_line": 367,
            "end_line": 400,
            "community": null
        },
        "node_id": "types\\message.py::12"
    },
    {
        "id": "types\\message.py::13",
        "input_type": "file",
        "content": "class Message(BaseModel):\n\n    @property\n    def parsed(self) -> Union[BaseModel, List[BaseModel]]:\n        \"\"\"Returns a list of all parsed content.\n\n        Example:\n            >>> class CustomModel(BaseModel):\n            ...     value: int\n            >>> parsed_content = CustomModel(value=42)\n            >>> message = Message(role=\"user\", content=[\"Text\", ContentBlock(parsed=parsed_content)])\n            >>> len(message.parsed)\n            1\n        \"\"\"\n        parsed_content = [c.parsed for c in self.content if c.parsed is not None]\n        return parsed_content[0] if len(parsed_content) == 1 else parsed_content",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 137,
            "span_ids": [
                "Message.parsed"
            ],
            "start_line": 402,
            "end_line": 415,
            "community": null
        },
        "node_id": "types\\message.py::13"
    },
    {
        "id": "types\\message.py::14",
        "input_type": "file",
        "content": "class Message(BaseModel):\n\n    def call_tools_and_collect_as_message(self, parallel=False, max_workers=None):\n        if parallel:\n            with ThreadPoolExecutor(max_workers=max_workers) as executor:\n                futures = [executor.submit(c.tool_call.call_and_collect_as_content_block) for c in self.content if c.tool_call]\n                content = [future.result() for future in as_completed(futures)]\n        else:\n            content = [c.tool_call.call_and_collect_as_content_block() for c in self.content if c.tool_call]\n        return Message(role=\"user\", content=content)",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 117,
            "span_ids": [
                "Message.call_tools_and_collect_as_message"
            ],
            "start_line": 417,
            "end_line": 424,
            "community": null
        },
        "node_id": "types\\message.py::14"
    },
    {
        "id": "types\\message.py::15",
        "input_type": "file",
        "content": "# HELPERS \ndef system(content: Union[AnyContent, List[AnyContent]]) -> Message:\n    \"\"\"\n    Create a system message with the given content.\n\n    Args:\n    content (str): The content of the system message.\n\n    Returns:\n    Message: A Message object with role set to 'system' and the provided content.\n    \"\"\"\n    return Message(role=\"system\", content=content)\n\n\ndef user(content: Union[AnyContent, List[AnyContent]]) -> Message:\n    \"\"\"\n    Create a user message with the given content.\n\n    Args:\n    content (str): The content of the user message.\n\n    Returns:\n    Message: A Message object with role set to 'user' and the provided content.\n    \"\"\"\n    return Message(role=\"user\", content=content)\n\n\ndef assistant(content: Union[AnyContent, List[AnyContent]]) -> Message:\n    \"\"\"\n    Create an assistant message with the given content.\n\n    Args:\n    content (str): The content of the assistant message.\n\n    Returns:\n    Message: A Message object with role set to 'assistant' and the provided content.\n    \"\"\"\n    return Message(role=\"assistant\", content=content)\n\n#XXX: Make a mixi for these properties.\ndef _content_to_text_only(content: List[ContentBlock]) -> str:\n    return _lstr(\"\\n\").join(\n            available_text\n            for c in content\n            if (available_text := (c.tool_result.text_only if c.tool_result else c.text))\n        )\n\n# Do we include the .text of a tool result? or its repr as in the current implementaiton?\n# What is the user using .text for? I just want to see the result of the tools. text_only should get us the text of the tool results; the tool_call_id is irrelevant.\ndef _content_to_text(content: List[ContentBlock]) -> str:\n    return _lstr(\"\\n\").join(\n            available_text\n            for c in content\n            if (available_text :=  c.text or repr(c.content))\n        )\n\n\n# want to enable a use case where the user can actually return a standrd oai chat format\n# This is a placehodler will likely come back later for this\nLMPParams = Dict[str, Any]\n# Well this is disappointing, I wanted to effectively type hint by doign that data sync meta, but eh, at elast we can still reference role or content this way. Probably wil lcan the dict sync meta. TypedDict is the ticket ell oh ell.\nMessageOrDict = Union[Message, Dict[str, str]]\n# Can support iamge prompts later.\nChat = List[\n    Message\n]  # [{\"role\": \"system\", \"content\": \"prompt\"}, {\"role\": \"user\", \"content\": \"message\"}]\nMultiTurnLMP = Callable[..., Chat]\nOneTurn = Callable[..., _lstr_generic]\n# This is the specific LMP that must accept history as an argument and can take any additional arguments\nChatLMP = Callable[[Chat, Any], Chat]\nLMP = Union[OneTurn, MultiTurnLMP, ChatLMP]\nInvocableLM = Callable[..., _lstr_generic]",
        "filepath": "src\\ell\\types\\message.py",
        "metadata": {
            "file_path": "src\\ell\\types\\message.py",
            "file_name": "message.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 652,
            "span_ids": [
                "_content_to_text",
                "impl:7",
                "system",
                "_content_to_text_only",
                "user",
                "Message.call_tools_and_collect_as_message",
                "assistant"
            ],
            "start_line": 426,
            "end_line": 500,
            "community": null
        },
        "node_id": "types\\message.py::15"
    },
    {
        "id": "types\\studio.py::1",
        "input_type": "file",
        "content": "from datetime import datetime, timezone\nimport enum\nfrom functools import cached_property\n\nimport sqlalchemy.types as types\n\nfrom ell.types.message import Any, Any, Field, Message, Optional\n\nfrom sqlmodel import Column, Field, SQLModel\nfrom typing import Optional\nfrom dataclasses import dataclass\nfrom typing import Dict, List, Literal, Union, Any, Optional\n\nfrom pydantic import BaseModel, field_validator\n\nfrom datetime import datetime\nfrom typing import Any, List, Optional\nfrom sqlmodel import Field, SQLModel, Relationship, JSON, Column\nfrom sqlalchemy import Index, func\n\nfrom typing import TypeVar, Any\n\ndef utc_now() -> datetime:\n    \"\"\"\n    Returns the current UTC timestamp.\n    Serializes to ISO-8601.\n    \"\"\"\n    return datetime.now(tz=timezone.utc)",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 166,
            "span_ids": [
                "utc_now",
                "imports"
            ],
            "start_line": 1,
            "end_line": 28,
            "community": null
        },
        "node_id": "types\\studio.py::1"
    },
    {
        "id": "types\\studio.py::2",
        "input_type": "file",
        "content": "class SerializedLMPUses(SQLModel, table=True):\n    \"\"\"\n    Represents the many-to-many relationship between SerializedLMPs.\n\n    This class is used to track which LMPs use or are used by other LMPs.\n    \"\"\"\n\n    lmp_user_id: Optional[str] = Field(default=None, foreign_key=\"serializedlmp.lmp_id\", primary_key=True, index=True)  # ID of the LMP that is being used\n    lmp_using_id: Optional[str] = Field(default=None, foreign_key=\"serializedlmp.lmp_id\", primary_key=True, index=True)  # ID of the LMP that is using the other LMP\r",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 138,
            "span_ids": [
                "SerializedLMPUses"
            ],
            "start_line": 31,
            "end_line": 39,
            "community": null
        },
        "node_id": "types\\studio.py::2"
    },
    {
        "id": "types\\studio.py::3",
        "input_type": "file",
        "content": "class UTCTimestamp(types.TypeDecorator[datetime]):\n    cache_ok = True\n    impl = types.TIMESTAMP\n    def process_result_value(self, value: datetime, dialect:Any):\n        return value.replace(tzinfo=timezone.utc)\n\n\ndef UTCTimestampField(index:bool=False, **kwargs:Any):\n    return Field(\n        sa_column=Column(UTCTimestamp(timezone=True), index=index, **kwargs))\n\n\nclass LMPType(str, enum.Enum):\n    LM = \"LM\"\n    TOOL = \"TOOL\"\n    MULTIMODAL = \"MULTIMODAL\"\n    OTHER = \"OTHER\"",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 125,
            "span_ids": [
                "UTCTimestampField",
                "UTCTimestamp",
                "UTCTimestamp.process_result_value",
                "LMPType"
            ],
            "start_line": 42,
            "end_line": 58,
            "community": null
        },
        "node_id": "types\\studio.py::3"
    },
    {
        "id": "types\\studio.py::4",
        "input_type": "file",
        "content": "class SerializedLMPBase(SQLModel):\n    lmp_id: Optional[str] = Field(default=None, primary_key=True)\n    name: str = Field(index=True)\n    source: str\n    dependencies: str\n    created_at: datetime = UTCTimestampField(index=True, nullable=False)\n\n    lmp_type: LMPType\n    api_params: Optional[Dict[str, Any]] = Field(default_factory=dict, sa_column=Column(JSON))\n    initial_free_vars: Optional[Dict[str, Any]] = Field(default_factory=dict, sa_column=Column(JSON))\n    initial_global_vars: Optional[Dict[str, Any]] = Field(default_factory=dict, sa_column=Column(JSON))\n    num_invocations: Optional[int] = Field(default=0)\n    commit_message: Optional[str] = Field(default=None)\n    version_number: Optional[int] = Field(default=None)",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 178,
            "span_ids": [
                "SerializedLMPBase"
            ],
            "start_line": 62,
            "end_line": 75,
            "community": null
        },
        "node_id": "types\\studio.py::4"
    },
    {
        "id": "types\\studio.py::5",
        "input_type": "file",
        "content": "class SerializedLMP(SerializedLMPBase, table=True):\n    invocations: List[\"Invocation\"] = Relationship(back_populates=\"lmp\")\n    used_by: Optional[List[\"SerializedLMP\"]] = Relationship(\n        back_populates=\"uses\",\n        link_model=SerializedLMPUses,\n        sa_relationship_kwargs=dict(\n            primaryjoin=\"SerializedLMP.lmp_id==SerializedLMPUses.lmp_user_id\",\n            secondaryjoin=\"SerializedLMP.lmp_id==SerializedLMPUses.lmp_using_id\",\n        ),\n    )\n    uses: List[\"SerializedLMP\"] = Relationship(\n        back_populates=\"used_by\",\n        link_model=SerializedLMPUses,\n        sa_relationship_kwargs=dict(\n            primaryjoin=\"SerializedLMP.lmp_id==SerializedLMPUses.lmp_using_id\",\n            secondaryjoin=\"SerializedLMP.lmp_id==SerializedLMPUses.lmp_user_id\",\n        ),\n    )\n\n    class Config:\n        table_name = \"serializedlmp\"\n        unique_together = [(\"version_number\", \"name\")]",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 214,
            "span_ids": [
                "SerializedLMP.Config",
                "SerializedLMP"
            ],
            "start_line": 78,
            "end_line": 99,
            "community": null
        },
        "node_id": "types\\studio.py::5"
    },
    {
        "id": "types\\studio.py::6",
        "input_type": "file",
        "content": "class InvocationTrace(SQLModel, table=True):\n    invocation_consumer_id: str = Field(foreign_key=\"invocation.id\", primary_key=True, index=True)\n    invocation_consuming_id: str = Field(foreign_key=\"invocation.id\", primary_key=True, index=True)\n\n# Should be subtyped for differnet kidns of LMPS.\n# XXX: Move all ofh te binary data out to a different table.\n# XXX: Need a flag that says dont store images.\n# XXX: Deprecate the args columns\nclass InvocationBase(SQLModel):\n    id: Optional[str] = Field(default=None, primary_key=True)\n    lmp_id: str = Field(foreign_key=\"serializedlmp.lmp_id\", index=True)\n    latency_ms: float\n    prompt_tokens: Optional[int] = Field(default=None)\n    completion_tokens: Optional[int] = Field(default=None)\n    state_cache_key: Optional[str] = Field(default=None)\n    created_at: datetime = UTCTimestampField(default=func.now(), nullable=False)\n    used_by_id: Optional[str] = Field(default=None, foreign_key=\"invocation.id\", index=True)\n    # global_vars and free_vars removed from here\r",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 249,
            "span_ids": [
                "InvocationTrace",
                "InvocationBase"
            ],
            "start_line": 101,
            "end_line": 118,
            "community": null
        },
        "node_id": "types\\studio.py::6"
    },
    {
        "id": "types\\studio.py::7",
        "input_type": "file",
        "content": "class InvocationContentsBase(SQLModel):\n    invocation_id: str = Field(foreign_key=\"invocation.id\", index=True, primary_key=True)\n    params: Optional[Dict[str, Any]] = Field(default=None, sa_column=Column(JSON))\n    results: Optional[Union[List[Message], Any]] = Field(default=None, sa_column=Column(JSON))\n    invocation_api_params: Optional[Dict[str, Any]] = Field(default=None, sa_column=Column(JSON))\n    global_vars: Optional[Dict[str, Any]] = Field(default=None, sa_column=Column(JSON))\n    free_vars: Optional[Dict[str, Any]] = Field(default=None, sa_column=Column(JSON))\n    is_external : bool = Field(default=False)",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 151,
            "span_ids": [
                "InvocationContentsBase"
            ],
            "start_line": 120,
            "end_line": 127,
            "community": null
        },
        "node_id": "types\\studio.py::7"
    },
    {
        "id": "types\\studio.py::8",
        "input_type": "file",
        "content": "class InvocationContentsBase(SQLModel):\n\n    @cached_property\n    def should_externalize(self) -> bool:\n        import json\n\n        json_fields = [\n            self.params,\n            self.results,\n            self.invocation_api_params,\n            self.global_vars,\n            self.free_vars\n        ]\n\n        total_size = sum(\n            len(json.dumps(field, default=(lambda x: json.dumps(x.model_dump(), default=str, ensure_ascii=False)\n                                           if isinstance(x, BaseModel) else str(x)), ensure_ascii=False).encode('utf-8'))\n            for field in json_fields if field is not None\n        )\n        # print(\"total_size\", total_size)\n\n        return total_size > 102400  # Precisely 100kb in bytes\n\nclass InvocationContents(InvocationContentsBase, table=True):\n    invocation: \"Invocation\" = Relationship(back_populates=\"contents\")",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 178,
            "span_ids": [
                "InvocationContents",
                "InvocationContentsBase.should_externalize"
            ],
            "start_line": 129,
            "end_line": 151,
            "community": null
        },
        "node_id": "types\\studio.py::8"
    },
    {
        "id": "types\\studio.py::9",
        "input_type": "file",
        "content": "class Invocation(InvocationBase, table=True):\n    lmp: SerializedLMP = Relationship(back_populates=\"invocations\")\n    consumed_by: List[\"Invocation\"] = Relationship(\n        back_populates=\"consumes\",\n        link_model=InvocationTrace,\n        sa_relationship_kwargs=dict(\n            primaryjoin=\"Invocation.id==InvocationTrace.invocation_consumer_id\",\n            secondaryjoin=\"Invocation.id==InvocationTrace.invocation_consuming_id\",\n        ),\n    )\n    consumes: List[\"Invocation\"] = Relationship(\n        back_populates=\"consumed_by\",\n        link_model=InvocationTrace,\n        sa_relationship_kwargs=dict(\n            primaryjoin=\"Invocation.id==InvocationTrace.invocation_consuming_id\",\n            secondaryjoin=\"Invocation.id==InvocationTrace.invocation_consumer_id\",\n        ),\n    )\n    used_by: Optional[\"Invocation\"] = Relationship(back_populates=\"uses\", sa_relationship_kwargs={\"remote_side\": \"Invocation.id\"})\n    uses: List[\"Invocation\"] = Relationship(back_populates=\"used_by\")\n    contents: InvocationContents = Relationship(back_populates=\"invocation\")\n\n    __table_args__ = (\n        Index('ix_invocation_lmp_id_created_at', 'lmp_id', 'created_at'),\n        Index('ix_invocation_created_at_latency_ms', 'created_at', 'latency_ms'),\n        Index('ix_invocation_created_at_tokens', 'created_at', 'prompt_tokens', 'completion_tokens'),\n    )",
        "filepath": "src\\ell\\types\\studio.py",
        "metadata": {
            "file_path": "src\\ell\\types\\studio.py",
            "file_name": "studio.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 284,
            "span_ids": [
                "Invocation"
            ],
            "start_line": 153,
            "end_line": 180,
            "community": null
        },
        "node_id": "types\\studio.py::9"
    },
    {
        "id": "util\\__init__.py::1",
        "input_type": "file",
        "content": "from .closure import lexically_closured_source",
        "filepath": "src\\ell\\util\\__init__.py",
        "metadata": {
            "file_path": "src\\ell\\util\\__init__.py",
            "file_name": "__init__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 10,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 1,
            "community": null
        },
        "node_id": "util\\__init__.py::1"
    },
    {
        "id": "util\\_warnings.py::1",
        "input_type": "file",
        "content": "from typing import Any, Optional\nfrom colorama import Fore, Style\n\nfrom ell.configurator import config\nimport logging\nlogger = logging.getLogger(__name__)\n\n\ndef _no_api_key_warning(model, client_to_use : Optional[Any], name = None,  long=False, error=False):\n    color = Fore.RED if error else Fore.LIGHTYELLOW_EX\n    prefix = \"ERROR\" if error else \"WARNING\"\n    # openai default\n    client_to_use_name = client_to_use.__class__.__name__ if (client_to_use) else \"OpenAI\"\n    client_to_use_module = client_to_use.__class__.__module__ if (client_to_use) else \"openai\"\n    lmp_name = f\"used by LMP `{name}` \" if name else \"\"\n    return f\"\"\"{color}{prefix}: No API key found for model `{model}` {lmp_name}using client `{client_to_use_name}`\"\"\" + (f\"\"\".\n\nTo fix this:\n* Set your API key in the appropriate environment variable for your chosen provider\n* Or, specify a client explicitly in the decorator:\n    ```\n    import ell\n    from {client_to_use_module} import {client_to_use_name}\n                                \n    @ell.simple(model=\"{model}\", client={client_to_use_name}(api_key=your_api_key))\n    def your_lmp_name(...):\n        ...\n    ```\n* Or explicitly specify the client when calling the LMP:\n\n    ```\n    your_lmp_name(..., client={client_to_use_name}(api_key=your_api_key))\n    ```\n\"\"\" if long else \" at time of definition. Can be okay if custom client specified later! https://docs.ell.so/core_concepts/models_and_api_clients.html \") + f\"{Style.RESET_ALL}\"",
        "filepath": "src\\ell\\util\\_warnings.py",
        "metadata": {
            "file_path": "src\\ell\\util\\_warnings.py",
            "file_name": "_warnings.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 365,
            "span_ids": [
                "_no_api_key_warning",
                "imports"
            ],
            "start_line": 1,
            "end_line": 34,
            "community": null
        },
        "node_id": "util\\_warnings.py::1"
    },
    {
        "id": "util\\_warnings.py::2",
        "input_type": "file",
        "content": "def _warnings(model, fn, default_client_from_decorator):\n\n        if not default_client_from_decorator:\n            # Check to see if the model is registered and warn the user we're gonna defualt to OpenAI.\n\n            if model not in config.registry:\n                logger.warning(f\"\"\"{Fore.LIGHTYELLOW_EX}WARNING: Model `{model}` is used by LMP `{fn.__name__}` but no client could be found that supports `{model}`. Defaulting to use the OpenAI client `{config.default_client}` for `{model}`. This is likely because you've spelled the model name incorrectly or are using a newer model from a provider added after this ell version was released. \n                            \n* If this is a mistake either specify a client explicitly in the decorator:\n```python\nimport ell\nell.simple(model, client=my_client)\ndef {fn.__name__}(...):\n    ...\n```\nor explicitly specify the client when the calling the LMP:\n\n```python\nell.simple(model, client=my_client)(...)\n```\n{Style.RESET_ALL}\"\"\")\n            elif (client_to_use := config.registry[model].default_client) is None or not client_to_use.api_key:\n                logger.warning(_no_api_key_warning(model, fn.__name__, client_to_use, long=False))\n\n\ndef _autocommit_warning():\n    if (config.get_client_for(\"gpt-4o-mini\")[0] is None):\n        logger.warning(f\"{Fore.LIGHTYELLOW_EX}WARNING: Autocommit is enabled but no OpenAI client found for autocommit model 'gpt-4o-mini' (set your OpenAI API key). Commit messages will not be written.{Style.RESET_ALL}\")\n        return True\n    return False",
        "filepath": "src\\ell\\util\\_warnings.py",
        "metadata": {
            "file_path": "src\\ell\\util\\_warnings.py",
            "file_name": "_warnings.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 353,
            "span_ids": [
                "_autocommit_warning",
                "_warnings"
            ],
            "start_line": 37,
            "end_line": 68,
            "community": null
        },
        "node_id": "util\\_warnings.py::2"
    },
    {
        "id": "util\\closure.py::1",
        "input_type": "file",
        "content": "\"\"\"\nThis should do the following.\n# prompt_consts.py\nimport math\ndef test():\n    return math.sin(10)\n\n# lol3.py\nimport prompt_consts\n\nX = 7\ndef xD():\n    print(X)\n    return prompt_consts.test()\n\n###\nOur goal is to use AST & dill to get a full lexical closured source of xD, with the exception of modules that are stored in site-packages. For example.\n\nlexical_extration(xD) returns\n#closure.py\nimport math\ndef test():\n    return math.sin(10)\n\nX = 7 \ndef xD():\n    print(X)\n    return test()\n\n\"\"\"\nimport collections\nimport ast\nimport hashlib\nimport itertools\nfrom typing import Any, Dict, Iterable, Optional, Set, Tuple, Callable\nimport dill\nimport inspect\nimport types\nfrom dill.source import getsource\nimport re\nfrom collections import deque\nimport black\n\nfrom ell.util.serialization import is_immutable_variable\nfrom ell.util.should_import import should_import\n\nDELIM = \"$$$$$$$$$$$$$$$$$$$$$$$$$\"\nFORBIDDEN_NAMES = [\"ell\", \"lstr\"]",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 235,
            "span_ids": [
                "docstring"
            ],
            "start_line": 1,
            "end_line": 48,
            "community": null
        },
        "node_id": "util\\closure.py::1"
    },
    {
        "id": "util\\closure.py::2",
        "input_type": "file",
        "content": "def lexical_closure(\n    func: Any,\n    already_closed: Set[int] = None,\n    initial_call: bool = False,\n    recursion_stack: list = None,\n    forced_dependencies: Optional[Dict[str, Any]] = None\n) -> Tuple[str, Tuple[str, str], Set[str]]:\n    \"\"\"\n    Generate a lexical closure for a given function or callable.\n\n    Args:\n        func: The function or callable to process.\n        already_closed: Set of already processed function hashes.\n        initial_call: Whether this is the initial call to the function.\n        recursion_stack: Stack to keep track of the recursion path.\n\n    Returns:\n        A tuple containing:\n        - The full source code of the closure\n        - A tuple of (function source, dependencies source)\n        - A set of function hashes that this closure uses\n    \"\"\"\n    already_closed = already_closed or set()\n    uses = set()\n    forced_dependencies = forced_dependencies or {}\n    recursion_stack = recursion_stack or []\n\n    if hash(func) in already_closed:\n        return \"\", (\"\", \"\"), set()\n\n    recursion_stack.append(getattr(func, '__qualname__', str(func)))\n\n    outer_ell_func = func\n    while hasattr(func, \"__ell_func__\"):\n        func = func.__ell_func__\n\n    source = getsource(func, lstrip=True)\n    already_closed.add(hash(func))\n\n    globals_and_frees = _get_globals_and_frees(func)\n    dependencies, imports, modules = _process_dependencies(func, globals_and_frees, already_closed, recursion_stack, uses)\n    for k,v in forced_dependencies.items():\n        # Todo: dictionary not necessary\n        _process_signature_dependency(v, dependencies, already_closed, recursion_stack, uses, k)\n\n    cur_src = _build_initial_source(imports, dependencies, source)\n\n    module_src = _process_modules(modules, cur_src, already_closed, recursion_stack, uses)\n\n    dirty_src = _build_final_source(imports, module_src, dependencies, source)\n    dirty_src_without_func = _build_final_source(imports, module_src, dependencies, \"\")\n\n    CLOSURE_SOURCE[hash(func)] = dirty_src\n\n    dsrc = _clean_src(dirty_src_without_func)\n\n    # Format the sorce and dsrc soruce using Black\n    source = _format_source(source)\n    dsrc = _format_source(dsrc)\n\n    fn_hash = _generate_function_hash(source, dsrc, func.__qualname__)\n\n    _update_ell_func(outer_ell_func, source, dsrc, globals_and_frees['globals'], globals_and_frees['frees'], fn_hash, uses)\n\n    return (dirty_src, (source, dsrc), ({outer_ell_func} if not initial_call and hasattr(outer_ell_func, \"__ell_func__\") else uses))",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 575,
            "span_ids": [
                "lexical_closure"
            ],
            "start_line": 50,
            "end_line": 114,
            "community": null
        },
        "node_id": "util\\closure.py::2"
    },
    {
        "id": "util\\closure.py::3",
        "input_type": "file",
        "content": "def _format_source(source: str) -> str:\n    \"\"\"Format the source code using Black.\"\"\"\n    try:\n        return black.format_str(source, mode=black.Mode())\n    except:\n        # If Black formatting fails, return the original source\n        return source\n\ndef _get_globals_and_frees(func: Callable) -> Dict[str, Dict]:\n    \"\"\"Get global and free variables for a function.\"\"\"\n    globals_dict = collections.OrderedDict(globalvars(func))\n    frees_dict = collections.OrderedDict(dill.detect.freevars(func))\n\n    if isinstance(func, type):\n        for name, method in collections.OrderedDict(func.__dict__).items():\n            if isinstance(method, (types.FunctionType, types.MethodType)):\n                globals_dict.update(collections.OrderedDict(dill.detect.globalvars(method)))\n                frees_dict.update(collections.OrderedDict(dill.detect.freevars(method)))\n\n    return {'globals': globals_dict, 'frees': frees_dict}",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 193,
            "span_ids": [
                "_get_globals_and_frees",
                "_format_source"
            ],
            "start_line": 117,
            "end_line": 136,
            "community": null
        },
        "node_id": "util\\closure.py::3"
    },
    {
        "id": "util\\closure.py::4",
        "input_type": "file",
        "content": "def _process_dependencies(func, globals_and_frees, already_closed, recursion_stack, uses):\n    \"\"\"Process function dependencies.\"\"\"\n    dependencies = []\n    modules = deque()\n    imports = []\n\n    if isinstance(func, (types.FunctionType, types.MethodType)):\n        _process_default_kwargs(func, dependencies, already_closed, recursion_stack, uses)\n\n    for var_name, var_value in itertools.chain(globals_and_frees['globals'].items(), globals_and_frees['frees'].items()):\n        _process_variable(var_name, var_value, dependencies, modules, imports, already_closed, recursion_stack, uses)\n\n    return dependencies, imports, modules",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 130,
            "span_ids": [
                "_process_dependencies"
            ],
            "start_line": 138,
            "end_line": 150,
            "community": null
        },
        "node_id": "util\\closure.py::4"
    },
    {
        "id": "util\\closure.py::5",
        "input_type": "file",
        "content": "def _process_default_kwargs(func, dependencies, already_closed, recursion_stack, uses):\n    \"\"\"Process default keyword arguments and annotations of a function.\"\"\"\n    ps = inspect.signature(func).parameters\n    for name, param in ps.items():\n        if param.default is not inspect.Parameter.empty:\n            _process_signature_dependency(param.default, dependencies, already_closed, recursion_stack, uses, name)\n        if param.annotation is not inspect.Parameter.empty:\n            _process_signature_dependency(param.annotation, dependencies, already_closed, recursion_stack, uses, f\"{name}_annotation\")\n    if func.__annotations__.get('return') is not None:\n        _process_signature_dependency(func.__annotations__['return'], dependencies, already_closed, recursion_stack, uses, \"return_annotation\")\n    # XXX: In order to properly analyze this we should walk the AST rather than inspexting the signature; e.g. Field is FieldInfo not Field.\n    # I don't care about the actual default at time of execution just the symbols required to statically reproduce the prompt.\r",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 210,
            "span_ids": [
                "_process_default_kwargs"
            ],
            "start_line": 151,
            "end_line": 162,
            "community": null
        },
        "node_id": "util\\closure.py::5"
    },
    {
        "id": "util\\closure.py::6",
        "input_type": "file",
        "content": "def _process_signature_dependency(val, dependencies, already_closed, recursion_stack, uses, name: Optional[str] = None):\n    # Todo: Build general cattr like utility for unstructuring python objects with hooks that keep track of state variables.\n    # Todo: break up closure into types and functions.\n    # XXX: This is not exhaustive, we should determine should import on all dependencies\n\n    if name not in FORBIDDEN_NAMES:\n        try:\n            dep = None\n            _uses = None\n            if isinstance(val, (types.FunctionType, types.MethodType)):\n                dep, _, _uses = lexical_closure(val, already_closed=already_closed, recursion_stack=recursion_stack.copy())\n            elif isinstance(val, (list, tuple, set)):\n                for item in val:\n                    _process_signature_dependency(item, dependencies, already_closed, recursion_stack, uses)\n            else:\n                val_class = val if isinstance(val, type) else val.__class__\n                try:\n                    is_builtin = (val_class.__module__ == \"builtins\" or val_class.__module__ == \"__builtins__\")\n                except:\n                    is_builtin = False\n\n                if not is_builtin:\n                    if should_import(val_class.__module__):\n                        dependencies.append(dill.source.getimport(val_class, alias=val_class.__name__))\n                    else:\n                        dep, _, _uses = lexical_closure(val_class, already_closed=already_closed, recursion_stack=recursion_stack.copy())\n\n            if dep: dependencies.append(dep)\n            if _uses: uses.update(_uses)\n        except Exception as e:\n            _raise_error(f\"Failed to capture the lexical closure of parameter or annotation {name}\", e, recursion_stack)",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 345,
            "span_ids": [
                "_process_signature_dependency"
            ],
            "start_line": 164,
            "end_line": 194,
            "community": null
        },
        "node_id": "util\\closure.py::6"
    },
    {
        "id": "util\\closure.py::7",
        "input_type": "file",
        "content": "def _process_variable(var_name, var_value, dependencies, modules, imports, already_closed, recursion_stack , uses):\n    \"\"\"Process a single variable.\"\"\"\n    try:\n        name = inspect.getmodule(var_value).__name__\n        if should_import(name):\n            imports.append(dill.source.getimport(var_value, alias=var_name))\n            return\n    except:\n        pass\n\n    if isinstance(var_value, (types.FunctionType, type, types.MethodType)):\n        _process_callable(var_name, var_value, dependencies, already_closed, recursion_stack, uses)\n    elif isinstance(var_value, types.ModuleType):\n        _process_module(var_name, var_value, modules, imports, uses)\n    elif isinstance(var_value, types.BuiltinFunctionType):\n        imports.append(dill.source.getimport(var_value, alias=var_name))\n    else:\n        _process_other_variable(var_name, var_value, dependencies, uses)",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 185,
            "span_ids": [
                "_process_variable"
            ],
            "start_line": 197,
            "end_line": 214,
            "community": null
        },
        "node_id": "util\\closure.py::7"
    },
    {
        "id": "util\\closure.py::8",
        "input_type": "file",
        "content": "def _process_callable(var_name, var_value, dependencies, already_closed, recursion_stack, uses):\n    \"\"\"Process a callable (function, method, or class).\"\"\"\n    try:\n        module_is_ell = 'ell' in inspect.getmodule(var_value).__name__\n    except:\n        module_is_ell = False\n\n    if var_name not in FORBIDDEN_NAMES and not module_is_ell:\n        try:\n            dep, _, _uses = lexical_closure(var_value, already_closed=already_closed, recursion_stack=recursion_stack.copy())\n            dependencies.append(dep)\n            uses.update(_uses)\n        except Exception as e:\n            _raise_error(f\"Failed to capture the lexical closure of global or free variable {var_name}\", e, recursion_stack)",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 155,
            "span_ids": [
                "_process_callable"
            ],
            "start_line": 216,
            "end_line": 229,
            "community": null
        },
        "node_id": "util\\closure.py::8"
    },
    {
        "id": "util\\closure.py::9",
        "input_type": "file",
        "content": "def _process_module(var_name, var_value, modules, imports, uses):\n    \"\"\"Process a module.\"\"\"\n    if should_import(var_value.__name__):\n        imports.append(dill.source.getimport(var_value, alias=var_name))\n    else:\n        modules.append((var_name, var_value))\n\ndef _process_other_variable(var_name, var_value, dependencies, uses):\n    \"\"\"Process variables that are not callables or modules.\"\"\"\n    if isinstance(var_value, str) and '\\n' in var_value:\n        dependencies.append(f\"{var_name} = '''{var_value}'''\")\n    elif is_immutable_variable(var_value):\n        dependencies.append(f\"#<BV>\\n{var_name} = {repr(var_value)}\\n#</BV>\")\n    else:\n        dependencies.append(f\"#<BmV>\\n{var_name} = <{type(var_value).__name__} object>\\n#</BmV>\")",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 191,
            "span_ids": [
                "_process_module",
                "_process_other_variable"
            ],
            "start_line": 231,
            "end_line": 245,
            "community": null
        },
        "node_id": "util\\closure.py::9"
    },
    {
        "id": "util\\closure.py::10",
        "input_type": "file",
        "content": "def _build_initial_source(imports, dependencies, source):\n    \"\"\"Build the initial source code.\"\"\"\n    return f\"{DELIM}\\n\" + f\"\\n{DELIM}\\n\".join(imports + dependencies + [source]) + f\"\\n{DELIM}\\n\"\n\ndef _process_modules(modules, cur_src, already_closed, recursion_stack, uses):\n    \"\"\"Process module dependencies.\"\"\"\n    reverse_module_src = deque()\n    while modules:\n        mname, mval = modules.popleft()\n        mdeps = []\n        attrs_to_extract = get_referenced_names(cur_src.replace(DELIM, \"\"), mname)\n        for attr in attrs_to_extract:\n            _process_module_attribute(mname, mval, attr, mdeps, modules, already_closed, recursion_stack, uses)\n\n        mdeps.insert(0, f\"# Extracted from module: {mname}\")\n        reverse_module_src.appendleft(\"\\n\".join(mdeps))\n\n        cur_src = _dereference_module_names(cur_src, mname, attrs_to_extract)\n\n    return list(reverse_module_src)",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 220,
            "span_ids": [
                "_build_initial_source",
                "_process_modules"
            ],
            "start_line": 247,
            "end_line": 266,
            "community": null
        },
        "node_id": "util\\closure.py::10"
    },
    {
        "id": "util\\closure.py::11",
        "input_type": "file",
        "content": "def _process_module_attribute(mname, mval, attr, mdeps, modules, already_closed, recursion_stack, uses):\n    \"\"\"Process a single attribute of a module.\"\"\"\n    val = getattr(mval, attr)\n    if isinstance(val, (types.FunctionType, type, types.MethodType)):\n        try:\n            dep, _, dep_uses = lexical_closure(val, already_closed=already_closed, recursion_stack=recursion_stack.copy())\n            mdeps.append(dep)\n            uses.update(dep_uses)\n        except Exception as e:\n            _raise_error(f\"Failed to capture the lexical closure of {mname}.{attr}\", e, recursion_stack)\n    elif isinstance(val, types.ModuleType):\n        modules.append((attr, val))\n    else:\n        mdeps.append(f\"{attr} = {repr(val)}\")",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 167,
            "span_ids": [
                "_process_module_attribute"
            ],
            "start_line": 268,
            "end_line": 281,
            "community": null
        },
        "node_id": "util\\closure.py::11"
    },
    {
        "id": "util\\closure.py::12",
        "input_type": "file",
        "content": "def _dereference_module_names(cur_src, mname, attrs_to_extract):\n    \"\"\"Dereference module names in the source code.\"\"\"\n    for attr in attrs_to_extract:\n        cur_src = cur_src.replace(f\"{mname}.{attr}\", attr)\n    return cur_src\n\ndef _build_final_source(imports, module_src, dependencies, source):\n    \"\"\"Build the final source code.\"\"\"\n    seperated_dependencies = sorted(imports) + sorted(module_src) + sorted(dependencies) + ([source] if source else [])\n    seperated_dependencies = list(dict.fromkeys(seperated_dependencies))\n    return DELIM + \"\\n\" + f\"\\n{DELIM}\\n\".join(seperated_dependencies) + \"\\n\" + DELIM + \"\\n\"\n\ndef _generate_function_hash(source, dsrc, qualname):\n    \"\"\"Generate a hash for the function.\"\"\"\n    return \"lmp-\" + hashlib.md5(\"\\n\".join((source, dsrc, qualname)).encode()).hexdigest()",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 205,
            "span_ids": [
                "_build_final_source",
                "_generate_function_hash",
                "_dereference_module_names"
            ],
            "start_line": 283,
            "end_line": 297,
            "community": null
        },
        "node_id": "util\\closure.py::12"
    },
    {
        "id": "util\\closure.py::13",
        "input_type": "file",
        "content": "def _update_ell_func(outer_ell_func, source, dsrc, globals_dict, frees_dict, fn_hash, uses):\n    \"\"\"Update the ell function attributes.\"\"\"\n    formatted_source = _format_source(source)\n    formatted_dsrc = _format_source(dsrc)\n\n    if hasattr(outer_ell_func, \"__ell_func__\"):\n\n        outer_ell_func.__ell_closure__ = (formatted_source, formatted_dsrc, globals_dict, frees_dict)\n        outer_ell_func.__ell_hash__ = fn_hash\n        outer_ell_func.__ell_uses__ = uses\n\ndef _raise_error(message, exception, recursion_stack):\n    \"\"\"Raise an error with detailed information.\"\"\"\n    error_msg = f\"{message}. Error: {str(exception)}\\n\"\n    error_msg += f\"Recursion stack: {' -> '.join(recursion_stack)}\"\n    # print(error_msg)\n    raise Exception(error_msg)",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 188,
            "span_ids": [
                "_update_ell_func",
                "_raise_error"
            ],
            "start_line": 299,
            "end_line": 315,
            "community": null
        },
        "node_id": "util\\closure.py::13"
    },
    {
        "id": "util\\closure.py::14",
        "input_type": "file",
        "content": "def get_referenced_names(code: str, module_name: str):\n    \"\"\"\n    This function takes a block of code and a module name as input. It parses the code into an Abstract Syntax Tree (AST)\n    and walks through the tree to find all instances where an attribute of the module is referenced in the code.\n\n    Parameters:\n    code (str): The block of code to be parsed.\n    module_name (str): The name of the module to look for in the code.\n\n    Returns:\n    list: A list of all attributes of the module that are referenced in the code.\n    \"\"\"\n    # Remove content between #<BV> and #</BV> tags\n    code = re.sub(r'#<BV>\\n.*?\\n#</BV>', '', code, flags=re.DOTALL)\n\n    # Remove content between #<BmV> and #</BmV> tags\n    code = re.sub(r'#<BmV>\\n.*?\\n#</BmV>', '', code, flags=re.DOTALL)\n\n    tree = ast.parse(code)\n    referenced_names = []\n\n    for node in ast.walk(tree):\n        if isinstance(node, ast.Attribute):\n            if isinstance(node.value, ast.Name) and node.value.id == module_name:\n                referenced_names.append(node.attr)\n\n    return referenced_names\n\nCLOSURE_SOURCE: Dict[str, str] = {}",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 285,
            "span_ids": [
                "get_referenced_names",
                "impl:5"
            ],
            "start_line": 318,
            "end_line": 346,
            "community": null
        },
        "node_id": "util\\closure.py::14"
    },
    {
        "id": "util\\closure.py::15",
        "input_type": "file",
        "content": "def lexically_closured_source(func, forced_dependencies: Optional[Dict[str, Any]] = None):\n    \"\"\"\n    Generate a lexically closured source for a given function.\n\n    This function takes a callable object (function, method, or class) and generates\n    a lexically closured source code. It captures all the dependencies, including\n    global variables, free variables, and nested functions, to create a self-contained\n    version of the function that can be executed independently.\n\n    Args:\n        func (Callable): The function or callable object to process.\n        forced_dependencies (Optional[Dict[str, Any]]): A dictionary of additional\n            dependencies to include in the closure. Keys are variable names, and\n            values are the corresponding objects.\n\n    Returns:\n        Tuple[str, Set[Any]]: A tuple containing two elements:\n            1. The lexically closured source code as a string.\n            2. A set of function objects that this closure uses.\n\n    Raises:\n        ValueError: If the input is not a callable object.\n\n    Example:\n        def outer(x):\n            y = 10\n            def inner():\n                return x + y\n            return inner\n\n        closured_source, uses = lexically_closured_source(outer)\n        print(closured_source)\n        # Output will include the source for both outer and inner functions,\n        # along with any necessary imports and variable definitions.\n\n    Note:\n        This function relies on the `lexical_closure` function to perform the\n        actual closure generation. It also uses the `__ell_closure__` attribute\n        of the function, which is expected to be set by the `lexical_closure` function.\n    \"\"\"\n    if not callable(func):\n        raise ValueError(\"Input must be a callable object (function, method, or class).\")\n    _, fnclosure, uses = lexical_closure(func, initial_call=True, recursion_stack=[], forced_dependencies=forced_dependencies)\n    return func.__ell_closure__, uses",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 406,
            "span_ids": [
                "lexically_closured_source"
            ],
            "start_line": 348,
            "end_line": 391,
            "community": null
        },
        "node_id": "util\\closure.py::15"
    },
    {
        "id": "util\\closure.py::16",
        "input_type": "file",
        "content": "import ast\n\ndef _clean_src(dirty_src):\n    # Now remove all duplicates and preserve order\n    split_by_setion = filter(lambda x: len(x.strip()) > 0, dirty_src.split(DELIM))\n\n    # Now we need to remove all the duplicates\n    split_by_setion = list(dict.fromkeys(split_by_setion))\n\n    # Now we need to concat all together\n    all_imports = []\n    final_src = \"\\n\".join(split_by_setion)\n    out_final_src = final_src[:]\n    for line in final_src.split(\"\\n\"):\n        if line.startswith(\"import\") or line.startswith(\"from\"):\n            all_imports.append(line)\n            out_final_src = out_final_src.replace(line, \"\")\n\n    all_imports = \"\\n\".join(sorted(all_imports))\n    final_src = all_imports + \"\\n\" + out_final_src\n\n    # now replace all \"\\n\\n\\n\" or longer with \"\\n\\n\"\n    final_src = re.sub(r\"\\n{3,}\", \"\\n\\n\", final_src)\n\n    return final_src",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 222,
            "span_ids": [
                "_clean_src",
                "impl:7"
            ],
            "start_line": 393,
            "end_line": 417,
            "community": null
        },
        "node_id": "util\\closure.py::16"
    },
    {
        "id": "util\\closure.py::17",
        "input_type": "file",
        "content": "def is_function_called(func_name, source_code):\n    \"\"\"\n    Check if a function is called in the given source code.\n\n    Parameters:\n    func_name (str): The name of the function to check.\n    source_code (str): The source code to check.\n\n    Returns:\n    bool: True if the function is called, False otherwise.\n    \"\"\"\n    # Parse the source code into an AST\n    tree = ast.parse(source_code)\n\n    # Walk through all the nodes in the AST\n    for node in ast.walk(tree):\n        # If the node is a function call\n        if isinstance(node, ast.Call):\n            # If the function being called is the function we're looking for\n            if isinstance(node.func, ast.Name) and node.func.id == func_name:\n                return True\n\n    # If we've gone through all the nodes and haven't found a call to the function, it's not called\n    return False\n\n#!/usr/bin/env python\n#\nfrom dill.detect import nestedglobals\nimport inspect",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 209,
            "span_ids": [
                "is_function_called",
                "impl:8"
            ],
            "start_line": 419,
            "end_line": 453,
            "community": null
        },
        "node_id": "util\\closure.py::17"
    },
    {
        "id": "util\\closure.py::18",
        "input_type": "file",
        "content": "def globalvars(func, recurse=True, builtin=False):\n    \"\"\"get objects defined in global scope that are referred to by func\n\n    return a dict of {name:object}\"\"\"\n    while hasattr(func, \"__ell_func__\"):\n        func = func.__ell_func__\n    if inspect.ismethod(func): func = func.__func__\n    while hasattr(func, \"__ell_func__\"):\n        func = func.__ell_func__\n    if inspect.isfunction(func):\n        globs = vars(inspect.getmodule(sum)).copy() if builtin else {}\n        # get references from within closure\n        orig_func, func = func, set()\n        for obj in orig_func.__closure__ or {}:\n            try:\n                cell_contents = obj.cell_contents\n            except ValueError: # cell is empty\n                pass\n            else:\n                _vars = globalvars(cell_contents, recurse, builtin) or {}\n                func.update(_vars) #XXX: (above) be wary of infinte recursion?\n                globs.update(_vars)\n        # get globals\n        globs.update(orig_func.__globals__ or {})\n        # get names of references\n        if not recurse:\n            func.update(orig_func.__code__.co_names)\n        else:\n            func.update(nestedglobals(orig_func.__code__))\n            # find globals for all entries of func\n            for key in func.copy(): #XXX: unnecessary...?\n                nested_func = globs.get(key)\n                if nested_func is orig_func:\n                   #func.remove(key) if key in func else None\n                    continue  #XXX: globalvars(func, False)?\n                func.update(globalvars(nested_func, True, builtin))\n    elif inspect.iscode(func):\n        globs = vars(inspect.getmodule(sum)).copy() if builtin else {}\n       #globs.update(globals())\n        if not recurse:\n            func = func.co_names # get names\n        else:\n            orig_func = func.co_name # to stop infinite recursion\n            func = set(nestedglobals(func))\n            # find globals for all entries of func\n            for key in func.copy(): #XXX: unnecessary...?\n                if key is orig_func:\n                   #func.remove(key) if key in func else None\n                    continue  #XXX: globalvars(func, False)?\n                nested_func = globs.get(key)\n                func.update(globalvars(nested_func, True, builtin))\n    # elif inspect.isclass(func):\n    # XXX: We need to get lexical closures of all the methods and attributes of the class.\\\n    # In the future we should exhaustively walk the AST here.\n    else:\n        return {}\n    #NOTE: if name not in __globals__, then we skip it...\n    return dict((name,globs[name]) for name in func if name in globs)",
        "filepath": "src\\ell\\util\\closure.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure.py",
            "file_name": "closure.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 573,
            "span_ids": [
                "globalvars"
            ],
            "start_line": 455,
            "end_line": 512,
            "community": null
        },
        "node_id": "util\\closure.py::18"
    },
    {
        "id": "util\\closure_util.py::1",
        "input_type": "file",
        "content": "import ast\nimport importlib\nimport os\nimport black\nfrom dill.detect import nestedglobals\n\nimport inspect\n\nimport inspect\n\n#!/usr/bin/env python\n#\ndef globalvars(func, recurse=True, builtin=False):\n    \"\"\"get objects defined in global scope that are referred to by func\n\n    return a dict of {name:object}\"\"\"\n    while hasattr(func, \"__ell_func__\"):\n        func = func.__ell_func__\n    if inspect.ismethod(func): func = func.__func__\n    while hasattr(func, \"__ell_func__\"):\n        func = func.__ell_func__\n    if inspect.isfunction(func):\n        globs = vars(inspect.getmodule(sum)).copy() if builtin else {}\n        # get references from within closure\n        orig_func, func = func, set()\n        for obj in orig_func.__closure__ or {}:\n            try:\n                cell_contents = obj.cell_contents\n            except ValueError: # cell is empty\n                pass\n            else:\n                _vars = globalvars(cell_contents, recurse, builtin) or {}\n                func.update(_vars) #XXX: (above) be wary of infinte recursion?\n                globs.update(_vars)\n        # get globals\n        globs.update(orig_func.__globals__ or {})\n        # get names of references\n        if not recurse:\n            func.update(orig_func.__code__.co_names)\n        else:\n            func.update(nestedglobals(orig_func.__code__))\n            # find globals for all entries of func\n            for key in func.copy(): #XXX: unnecessary...?\n                nested_func = globs.get(key)\n                if nested_func is orig_func:\n                   #func.remove(key) if key in func else None\n                    continue  #XXX: globalvars(func, False)?\n                func.update(globalvars(nested_func, True, builtin))\n    elif inspect.iscode(func):\n        globs = vars(inspect.getmodule(sum)).copy() if builtin else {}\n       #globs.update(globals())\n        if not recurse:\n            func = func.co_names # get names\n        else:\n            orig_func = func.co_name # to stop infinite recursion\n            func = set(nestedglobals(func))\n            # find globals for all entries of func\n            for key in func.copy(): #XXX: unnecessary...?\n                if key is orig_func:\n                   #func.remove(key) if key in func else None\n                    continue  #XXX: globalvars(func, False)?\n                nested_func = globs.get(key)\n                func.update(globalvars(nested_func, True, builtin))\n    else:\n        return {}\n    #NOTE: if name not in __globals__, then we skip it...\n    return dict((name,globs[name]) for name in func if name in globs)",
        "filepath": "src\\ell\\util\\closure_util.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure_util.py",
            "file_name": "closure_util.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 564,
            "span_ids": [
                "globalvars",
                "imports"
            ],
            "start_line": 1,
            "end_line": 75,
            "community": null
        },
        "node_id": "util\\closure_util.py::1"
    },
    {
        "id": "util\\closure_util.py::2",
        "input_type": "file",
        "content": "def is_function_called(func_name, source_code):\n    \"\"\"\n    Check if a function is called in the given source code.\n\n    Parameters:\n    func_name (str): The name of the function to check.\n    source_code (str): The source code to check.\n\n    Returns:\n    bool: True if the function is called, False otherwise.\n    \"\"\"\n    # Parse the source code into an AST\n    tree = ast.parse(source_code)\n\n    # Walk through all the nodes in the AST\n    for node in ast.walk(tree):\n        # If the node is a function call\n        if isinstance(node, ast.Call):\n            # If the function being called is the function we're looking for\n            if isinstance(node.func, ast.Name) and node.func.id == func_name:\n                return True\n\n    # If we've gone through all the nodes and haven't found a call to the function, it's not called\n    return False",
        "filepath": "src\\ell\\util\\closure_util.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure_util.py",
            "file_name": "closure_util.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 191,
            "span_ids": [
                "is_function_called"
            ],
            "start_line": 78,
            "end_line": 101,
            "community": null
        },
        "node_id": "util\\closure_util.py::2"
    },
    {
        "id": "util\\closure_util.py::3",
        "input_type": "file",
        "content": "def get_referenced_names(code: str, module_name: str):\n    \"\"\"\n    This function takes a block of code and a module name as input. It parses the code into an Abstract Syntax Tree (AST)\n    and walks through the tree to find all instances where an attribute of the module is referenced in the code.\n\n    Parameters:\n    code (str): The block of code to be parsed.\n    module_name (str): The name of the module to look for in the code.\n\n    Returns:\n    list: A list of all attributes of the module that are referenced in the code.\n    \"\"\"\n    tree = ast.parse(code)\n    referenced_names = []\n\n    for node in ast.walk(tree):\n        if isinstance(node, ast.Attribute):\n            if isinstance(node.value, ast.Name) and node.value.id == module_name:\n                referenced_names.append(node.attr)\n\n    return referenced_names",
        "filepath": "src\\ell\\util\\closure_util.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure_util.py",
            "file_name": "closure_util.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 178,
            "span_ids": [
                "get_referenced_names"
            ],
            "start_line": 104,
            "end_line": 124,
            "community": null
        },
        "node_id": "util\\closure_util.py::3"
    },
    {
        "id": "util\\closure_util.py::4",
        "input_type": "file",
        "content": "def should_import(module_name : str):\n    \"\"\"\n    This function checks if a module should be imported based on its origin.\n    It returns False if the module is in the local directory or if the module's spec is None.\n    Otherwise, it returns True.\n\n    Returns:\n    bool: True if the module should be imported, False otherwise.\n    \"\"\"\n\n    # Define the local directory\n    DIRECTORY_TO_WATCH = os.environ.get(\"DIRECTORY_TO_WATCH\", os.getcwd())\n\n    # Get the module's spec\n    spec = importlib.util.find_spec(module_name)\n\n    if module_name.startswith(\"ell\"):\n        return True\n\n    # Return False if the spec is None or if the spec's origin starts with the local directory\n    if spec is None or (spec.origin is not None and spec.origin.startswith(DIRECTORY_TO_WATCH)):\n        return False\n\n    # Otherwise, return True\n    return True\n\n\ndef format_source(source: str) -> str:\n    \"\"\"Format the source code using Black.\"\"\"\n    try:\n        return black.format_str(source, mode=black.Mode())\n    except:\n        # If Black formatting fails, return the original source\n        return source",
        "filepath": "src\\ell\\util\\closure_util.py",
        "metadata": {
            "file_path": "src\\ell\\util\\closure_util.py",
            "file_name": "closure_util.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 236,
            "span_ids": [
                "should_import",
                "format_source"
            ],
            "start_line": 127,
            "end_line": 160,
            "community": null
        },
        "node_id": "util\\closure_util.py::4"
    },
    {
        "id": "util\\differ.py::1",
        "input_type": "file",
        "content": "from ell.configurator import config\nfrom ell.lmp.simple import simple\nimport difflib\n\n# Todo: update this for single change stuff so that it doesn't summarize small chage but says it specifically.\n@simple(config.autocommit_model, temperature=0.2, exempt_from_tracking=True, max_tokens=500)\ndef write_commit_message_for_diff(old : str, new : str) -> str:\n    \"\"\"You are an expert programmer whose goal is to write commit messages based on diffs.\nYou will be given two version of source code and their unified diff.\nYou will be expected to write a commit message that describes the changes between the two versions.\n\nFollow these guidelines:\n1. Your commit message should be at most one sentence and highly specific to the changes made. Don't just discuss the functions changed but how they were specifically changed.\n2. Your commit message cannot be more than 10 words so use sentence fragments and be concise.\n3. The @ell.simple decorator turns a function into a call to a language model program: the function's docstring is the system prompt and the string returned is the user prompt. \n4. It is extremely important that if the system prompt or user prompt changes, your commit message must say what specifically changed, rather than vaguely saying they were updated or changed.\n5. It is extremely important that you never refer to a @ell.simple docstring as a docstring: it is a system prompt. \n6. Do NOT say why a change was done, say what specifically changed.\n7. Consider all changes ot the program including the globals and free variables\n\nExample response:\n'''\nUpdate model temperature and refine system prompt wording:\n* Changed temperature from 0.5 to 0.7.\n* Updated \"with specificity, brevity, and good grammar\" to \"clearly and concisely\" in system prompt.\n* The `questions` param was assigned type List[Question]\n'''\nResponse format:\n<Short commit message summarizing all the changes with specificity>:\n* <Bulleted list of each specific change>.\n\"\"\"\n    clean_program_of_all_bv_tags = lambda program : program.replace(\"#<BV>\", \"\").replace(\"#</BV>\", \"\").replace(\"#<BmV>\", \"\").replace(\"#</BmV>\", \"\")\n    old_clean = clean_program_of_all_bv_tags(old)\n    new_clean = clean_program_of_all_bv_tags(new)\n\n    diff = \"\\n\".join(difflib.unified_diff(old_clean.splitlines(), new_clean.splitlines(), lineterm=''))\n\n    return f\"\"\"Write a commit message succinctly and specifically describing the changes between these two versions of a program.\nOld version:\n```\n{old_clean}\n```\n\nNew version:\n```\n{new_clean}\n```\n\nUnified diff:\n{diff}\n\"\"\"",
        "filepath": "src\\ell\\util\\differ.py",
        "metadata": {
            "file_path": "src\\ell\\util\\differ.py",
            "file_name": "differ.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 569,
            "span_ids": [
                "write_commit_message_for_diff",
                "imports"
            ],
            "start_line": 1,
            "end_line": 51,
            "community": null
        },
        "node_id": "util\\differ.py::1"
    },
    {
        "id": "util\\differ.py::2",
        "input_type": "file",
        "content": "if __name__ == \"__main__\":\n\n    from ell.configurator import config\n    config.verbose = True\n\n    test_version_1 = '''import ell\nimport numpy as np\n\n@ell.simple(model=\"gpt-4o-mini\")\ndef come_up_with_a_premise_for_a_joke_about(topic : str):\n    \"\"\"You are an incredibly funny comedian. Come up with a premise for a joke about topic\"\"\"\n    return f\"come up with a premise for a joke about {topic}\"\n\ndef get_random_length():\n    return int(np.random.beta(2, 5) * 300)\n\n@ell.simple(model=\"gpt-4o-mini\")\ndef joke(topic : str):\n    \"\"\"You are a funny comedian. You respond in scripts for a standup comedy skit.\"\"\"\n    return f\"Act out a full joke. Make your script {get_random_length()} words long. Here's the premise: {come_up_with_a_premise_for_a_joke_about(topic)}\"'''\n\n    test_version_2 = '''import ell\nimport numpy as np\n\n@ell.simple(model=\"gpt-4o-mini\")\ndef come_up_with_a_premise_for_a_joke_about(topic : str):\n    \"\"\"You are an incredibly funny comedian. Come up with a premise for a joke about topic\"\"\"\n    return f\"come up with a premise for a joke about {topic}\"\n\ndef get_random_length():\n    return int(np.random.beta(2, 5) * 300)\n\n@ell.simple(model=\"gpt-4o-mini\")\ndef joke(topic : str):\n    \"\"\"You are a funny comedian. You respond in scripts for skits.\"\"\"\n    return f\"Act out a full joke. Make your script {get_random_length()} words long. Here's the premise: {come_up_with_a_premise_for_a_joke_about(topic)}\"'''\n\n    (write_commit_message_for_diff(test_version_1, test_version_2))",
        "filepath": "src\\ell\\util\\differ.py",
        "metadata": {
            "file_path": "src\\ell\\util\\differ.py",
            "file_name": "differ.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 404,
            "span_ids": [
                "impl"
            ],
            "start_line": 53,
            "end_line": 90,
            "community": null
        },
        "node_id": "util\\differ.py::2"
    },
    {
        "id": "util\\plot_ascii.py::1",
        "input_type": "file",
        "content": "import sys\nfrom PIL import Image, ImageDraw, ImageFont\nimport numpy as np\nimport logging\nimport os\n\n# Load pre-rendered character bitmaps\ntry:\n    package_dir = os.path.dirname(__file__)\n    bitmaps_path = os.path.join(package_dir, 'char_bitmaps.npy')\n    data = np.load(bitmaps_path, allow_pickle=True).item()\n    char_bitmaps = data['char_bitmaps']\n    max_char_width = data['max_char_width']\n    max_char_height = data['max_char_height']\n\n\n    ASCII_CHARS = \" .:-=+*#%@\"\n    def plot_ascii(\n        image: Image.Image,\n        width: int = 100,\n        color: bool = True,\n    ):\n        \"\"\"\n        Convert a PIL Image to ASCII art using pre-rendered character bitmaps and print it to the console with optional coloring.\n        \"\"\"\n        num_chars = len(ASCII_CHARS)\n\n\n        # Adjust the scaling factor to compensate for character aspect ratio\n        scale = 0.5  # You can tweak this value based on your terminal's character dimensions\n        aspect_ratio = image.height / image.width\n        new_width = width\n        new_height = int(aspect_ratio * new_width * (max_char_height / max_char_width) * scale)\n        image = image.resize((new_width * max_char_width, new_height * max_char_height)).convert('RGB')\n\n        # Convert image to NumPy array\n        img_array = np.array(image)\n\n        # Compute brightness using luminance formula\n        luminance = 0.2126 * img_array[:, :, 0] + 0.7152 * img_array[:, :, 1] + 0.0722 * img_array[:, :, 2]\n\n        # Normalize brightness to range 0-1\n        brightness_normalized = luminance / 255\n\n        if color:\n            # Get RGB values for coloring\n            r = img_array[:, :, 0]\n            g = img_array[:, :, 1]\n            b = img_array[:, :, 2]\n\n        # Compute the number of blocks\n        y_blocks = new_height\n        x_blocks = new_width\n\n        # Reshape brightness_normalized to (y_blocks, max_char_height, x_blocks, max_char_width)\n        brightness_blocks = brightness_normalized.reshape(y_blocks, max_char_height, x_blocks, max_char_width)\n        brightness_blocks = brightness_blocks.mean(axis=(1, 3))  # Average over each block\n\n        # Normalize again if necessary\n        brightness_blocks = brightness_blocks / brightness_blocks.max()\n\n        # Vectorize the selection of ASCII characters\n        indices = np.digitize(brightness_blocks, np.linspace(0, 1, num_chars)) - 1\n        indices = np.clip(indices, 0, num_chars - 1)\n        ascii_chars = np.array(list(ASCII_CHARS))[indices]\n\n        if color:\n            # Compute average color for each block\n            r_blocks = r.reshape(y_blocks, max_char_height, x_blocks, max_char_width).mean(axis=(1, 3)).astype(int)\n            g_blocks = g.reshape(y_blocks, max_char_height, x_blocks, max_char_width).mean(axis=(1, 3)).astype(int)\n            b_blocks = b.reshape(y_blocks, max_char_height, x_blocks, max_char_width).mean(axis=(1, 3)).astype(int)\n\n            # Convert RGB to 8-bit color code\n            color_codes = 16 + (36 * (r_blocks // 51)) + (6 * (g_blocks // 51)) + (b_blocks // 51)\n            color_codes = color_codes.astype(str)\n\n            # Create colored ASCII characters\n            colored_ascii = np.char.add(np.char.add(\"\\033[38;5;\", color_codes), \"m\")\n            colored_ascii = np.char.add(colored_ascii, np.char.add(ascii_chars, \"\\033[0m\"))\n\n            # Join characters into lines\n            ascii_image = [\"\".join(row) for row in colored_ascii]\n        else:\n            ascii_image = [\"\".join(row) for row in ascii_chars]\n\n        # Print the ASCII image\n        return ascii_image\nexcept FileNotFoundError:\n    def plot_ascii(\n        image: Image.Image,\n        width: int = 100,\n        color: bool = True,\n    ):\n        return \"<image>\"",
        "filepath": "src\\ell\\util\\plot_ascii.py",
        "metadata": {
            "file_path": "src\\ell\\util\\plot_ascii.py",
            "file_name": "plot_ascii.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 896,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 94,
            "community": null
        },
        "node_id": "util\\plot_ascii.py::1"
    },
    {
        "id": "util\\plot_ascii.py::2",
        "input_type": "file",
        "content": "# For packaging .\ndef render_and_save_char_bitmaps(\n    font_path: str = \"Courier New.ttf\",\n    font_size: int = 10,\n    output_path: str = \"char_bitmaps.npy\"\n) -> None:\n    num_chars = len(ASCII_CHARS)\n    char_bitmaps = {}\n    max_char_width, max_char_height = 0, 0\n\n    try:\n        font = ImageFont.truetype(font_path, font_size)\n    except IOError:\n        logging.error(f\"Font file '{font_path}' not found. Please provide a valid path.\")\n        sys.exit(1)\n\n    for char in ASCII_CHARS:\n        char_image = Image.new('L', (font_size, font_size), color=255)\n        draw = ImageDraw.Draw(char_image)\n        bbox = font.getbbox(char)\n        w = bbox[2] - bbox[0]\n        h = bbox[3] - bbox[1]\n        draw.text(((font_size - w) / 2, (font_size - h) / 2), char, fill=0, font=font)\n        bitmap = np.array(char_image) / 255  # Normalize to 0-1\n        char_bitmaps[char] = bitmap\n        max_char_width = max(max_char_width, bitmap.shape[1])\n        max_char_height = max(max_char_height, bitmap.shape[0])\n\n    # Save the character bitmaps and max dimensions\n    data = {\n        'char_bitmaps': char_bitmaps,\n        'max_char_width': max_char_width,\n        'max_char_height': max_char_height\n    }\n    np.save(output_path, data)\n    logging.info(f\"Character bitmaps saved to '{output_path}'.\")",
        "filepath": "src\\ell\\util\\plot_ascii.py",
        "metadata": {
            "file_path": "src\\ell\\util\\plot_ascii.py",
            "file_name": "plot_ascii.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 354,
            "span_ids": [
                "render_and_save_char_bitmaps",
                "imports"
            ],
            "start_line": 96,
            "end_line": 133,
            "community": null
        },
        "node_id": "util\\plot_ascii.py::2"
    },
    {
        "id": "util\\serialization.py::1",
        "input_type": "file",
        "content": "# Global converter\nimport base64\nimport hashlib\nfrom io import BytesIO\nimport json\nimport cattrs\nimport numpy as np\nfrom pydantic import BaseModel\nimport PIL\nfrom ell.types._lstr import _lstr\n\n\npydantic_ltype_aware_cattr = cattrs.Converter()\n\ndef serialize_image(img):\n    buffer = BytesIO()\n    img.save(buffer, format=\"PNG\")\n    return \"data:image/png;base64,\" + base64.b64encode(buffer.getvalue()).decode()\n\n\n# Register hooks for complex types\npydantic_ltype_aware_cattr.register_unstructure_hook(\n    np.ndarray,\n    lambda arr: {\n        \"content\": serialize_image(PIL.Image.fromarray(arr)),\n        \"__limage\": True\n    } if arr.ndim == 3 else (\n        {\n            \"content\": base64.b64encode(arr.tobytes()).decode(),\n            \"dtype\": str(arr.dtype),\n            \"shape\": arr.shape,\n            \"__lndarray\": True\n        }\n    )\n)\npydantic_ltype_aware_cattr.register_unstructure_hook(\n    set,\n    lambda s: list(sorted(s))\n)\npydantic_ltype_aware_cattr.register_unstructure_hook(\n    frozenset,\n    lambda s: list(sorted(s))\n)\n\n\npydantic_ltype_aware_cattr.register_unstructure_hook(\n    PIL.Image.Image,\n    lambda obj: {\n        \"content\": serialize_image(obj),\n        \"__limage\": True\n    }\n)\n\ndef unstructure_lstr(obj):\n    return dict(content=str(obj), **obj.__dict__, __lstr=True)\n\npydantic_ltype_aware_cattr.register_unstructure_hook(\n    _lstr,\n    unstructure_lstr\n)\n\npydantic_ltype_aware_cattr.register_unstructure_hook(\n    BaseModel,\n    lambda obj: obj.model_dump(exclude_none=True, exclude_unset=True)\n)",
        "filepath": "src\\ell\\util\\serialization.py",
        "metadata": {
            "file_path": "src\\ell\\util\\serialization.py",
            "file_name": "serialization.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 391,
            "span_ids": [
                "docstring",
                "impl:7",
                "unstructure_lstr",
                "serialize_image",
                "impl:3"
            ],
            "start_line": 2,
            "end_line": 66,
            "community": null
        },
        "node_id": "util\\serialization.py::1"
    },
    {
        "id": "util\\serialization.py::2",
        "input_type": "file",
        "content": "def get_immutable_vars(vars_dict):\n    converter = cattrs.Converter()\n\n    def handle_complex_types(obj):\n        if isinstance(obj, (int, float, str, bool, type(None))):\n            return obj\n        elif isinstance(obj, (list, tuple)):\n            return [handle_complex_types(item) if not isinstance(item, (int, float, str, bool, type(None))) else item for item in obj]\n        elif isinstance(obj, dict):\n            return {k: handle_complex_types(v) if not isinstance(v, (int, float, str, bool, type(None))) else v for k, v in obj.items()}\n        elif isinstance(obj, (set, frozenset)):\n            return list(sorted(handle_complex_types(item) if not isinstance(item, (int, float, str, bool, type(None))) else item for item in obj))\n        elif isinstance(obj, np.ndarray):\n            return obj.tolist()\n        else:\n            return f\"<Object of type {type(obj).__name__}>\"\n\n    converter.register_unstructure_hook(object, handle_complex_types)\n    x = converter.unstructure(vars_dict)\n    return x",
        "filepath": "src\\ell\\util\\serialization.py",
        "metadata": {
            "file_path": "src\\ell\\util\\serialization.py",
            "file_name": "serialization.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 230,
            "span_ids": [
                "get_immutable_vars"
            ],
            "start_line": 71,
            "end_line": 90,
            "community": null
        },
        "node_id": "util\\serialization.py::2"
    },
    {
        "id": "util\\serialization.py::3",
        "input_type": "file",
        "content": "def compute_state_cache_key(ipstr, fn_closure):\n    _global_free_vars_str = f\"{json.dumps(get_immutable_vars(fn_closure[2]), sort_keys=True, default=repr, ensure_ascii=False)}\"\n    _free_vars_str = f\"{json.dumps(get_immutable_vars(fn_closure[3]), sort_keys=True, default=repr, ensure_ascii=False)}\"\n    state_cache_key = hashlib.sha256(f\"{ipstr}{_global_free_vars_str}{_free_vars_str}\".encode('utf-8')).hexdigest()\n    return state_cache_key",
        "filepath": "src\\ell\\util\\serialization.py",
        "metadata": {
            "file_path": "src\\ell\\util\\serialization.py",
            "file_name": "serialization.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 111,
            "span_ids": [
                "compute_state_cache_key"
            ],
            "start_line": 93,
            "end_line": 97,
            "community": null
        },
        "node_id": "util\\serialization.py::3"
    },
    {
        "id": "util\\serialization.py::4",
        "input_type": "file",
        "content": "def prepare_invocation_params(params):\n    invocation_params = params\n\n    cleaned_invocation_params = pydantic_ltype_aware_cattr.unstructure(invocation_params)\n\n    # Thisis because we wneed the caching to work on the hash of a cleaned and serialized object.\n    jstr = json.dumps(cleaned_invocation_params, sort_keys=True, default=repr, ensure_ascii=False)\n\n    consumes = set()\n    import re\n    # XXX: Better than registering a hook in cattrs.\n    pattern = r'\"__origin_trace__\":\\s*\"frozenset\\({(.+?)}\\)\"'\n\n    # Find all matches in the jstr\n    matches = re.findall(pattern, jstr)\n\n    # Process each match and add to consumes set\n    for match in matches:\n        # Remove quotes and spaces, then split by comma\n        items = [item.strip().strip(\"'\") for item in match.split(',')]\n        consumes.update(items)\n    consumes = list(consumes)\n    # XXX: Only need to reload because of 'input' caching., we could skip this by making ultimate model caching rather than input hash caching; if prompt same use the same output.. irrespective of version.\n    return json.loads(jstr), jstr, consumes",
        "filepath": "src\\ell\\util\\serialization.py",
        "metadata": {
            "file_path": "src\\ell\\util\\serialization.py",
            "file_name": "serialization.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 260,
            "span_ids": [
                "prepare_invocation_params"
            ],
            "start_line": 100,
            "end_line": 123,
            "community": null
        },
        "node_id": "util\\serialization.py::4"
    },
    {
        "id": "util\\serialization.py::5",
        "input_type": "file",
        "content": "def is_immutable_variable(value):\n    \"\"\"\n    Check if a value is immutable.\n\n    This function determines whether the given value is of an immutable type in Python.\n    Immutable types are objects whose state cannot be modified after they are created.\n\n    Args:\n        value: Any Python object to check for immutability.\n\n    Returns:\n        bool: True if the value is immutable, False otherwise.\n\n    Note:\n        - This function checks for common immutable types in Python.\n        - Custom classes are considered mutable unless they explicitly implement\n          immutability (which this function doesn't check for).\n        - For some types like tuple, immutability is shallow (i.e., the tuple itself\n          is immutable, but its contents might not be).\n    \"\"\"\n    immutable_types = (\n        int, float, complex, str, bytes,\n        tuple, frozenset, type(None),\n        bool,  # booleans are immutable\n        range,  # range objects are immutable\n        slice,  # slice objects are immutable\n    )\n\n    if isinstance(value, immutable_types):\n        return True\n\n    # Check for immutable instances of mutable types\n    if isinstance(value, (tuple, frozenset)):\n        return all(is_immutable_variable(item) for item in value)\n\n    return False",
        "filepath": "src\\ell\\util\\serialization.py",
        "metadata": {
            "file_path": "src\\ell\\util\\serialization.py",
            "file_name": "serialization.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 263,
            "span_ids": [
                "is_immutable_variable"
            ],
            "start_line": 126,
            "end_line": 161,
            "community": null
        },
        "node_id": "util\\serialization.py::5"
    },
    {
        "id": "util\\should_import.py::1",
        "input_type": "file",
        "content": "import importlib.util\nimport os\nimport site\nimport sys\nimport sysconfig\nfrom pathlib import Path\n\n\ndef should_import(module_name: str, raise_on_error: bool = False) -> bool:\n    \"\"\"\n    Determines whether a module should be imported based on its origin.\n    Excludes local modules and standard library modules.\n\n    Args:\n        module_name (str): The name of the module to check.\n\n    Returns:\n        bool: True if the module should be imported (i.e., it's a third-party module), False otherwise.\n    \"\"\"\n    if module_name.startswith(\"ell\"):\n        return True\n    try:\n        try:\n            spec = importlib.util.find_spec(module_name)\n        except ValueError:\n            return False\n        if spec is None:\n            return False\n\n        origin = spec.origin\n        if origin is None:\n            return False\n        if spec.has_location:\n            origin_path = Path(origin).resolve()\n\n            site_packages = list(site.getsitepackages()) + (list(site.getusersitepackages()) if isinstance(site.getusersitepackages(), list) else [site.getusersitepackages()])\n\n            additional_paths = [Path(p).resolve() for p in sys.path if Path(p).resolve() not in map(Path, site_packages)]\n\n            project_root = Path(os.environ.get(\"ELL_PROJECT_ROOT\", os.getcwd())).resolve()\n\n            site_packages_paths = [Path(p).resolve() for p in site_packages]\n            stdlib_path = sysconfig.get_paths().get(\"stdlib\")\n            if stdlib_path:\n                site_packages_paths.append(Path(stdlib_path).resolve())\n\n            additional_paths = [Path(p).resolve() for p in additional_paths]\n            local_paths = [project_root]\n\n            cwd = Path.cwd().resolve()\n            additional_paths = [path for path in additional_paths if path != cwd]\n\n            for pkg in site_packages_paths:\n                if origin_path.is_relative_to(pkg):\n                    return True\n\n            for path in additional_paths:\n                if origin_path.is_relative_to(path):\n                    return False\n\n            for local in local_paths:\n                if origin_path.is_relative_to(local):\n                    return False\n\n        return True\n\n    except Exception as e:\n        if raise_on_error:\n            raise e\n        return True",
        "filepath": "src\\ell\\util\\should_import.py",
        "metadata": {
            "file_path": "src\\ell\\util\\should_import.py",
            "file_name": "should_import.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 457,
            "span_ids": [
                "should_import",
                "imports"
            ],
            "start_line": 1,
            "end_line": 70,
            "community": null
        },
        "node_id": "util\\should_import.py::1"
    },
    {
        "id": "util\\verbosity.py::1",
        "input_type": "file",
        "content": "import sys\nimport hashlib\nimport shutil\nimport textwrap\nfrom typing import Dict, Tuple, List, Any, Optional\nfrom colorama import Fore, Style, init\nfrom ell.types import Message\nfrom ell.configurator import config\nimport logging\nfrom functools import lru_cache\nimport threading\nfrom ell.types.message import LMP, ContentBlock, _content_to_text\nimport requests\n\nfrom ell.util.plot_ascii import plot_ascii\n\n# Initialize colorama\ninit(autoreset=True)\n\n# Define colors and styles\nELL_COLORS = {k: v for k, v in vars(Fore).items() if k not in ['RESET', 'BLACK', 'LIGHTBLACK_EX']}\nBOLD = Style.BRIGHT\nUNDERLINE = '\\033[4m'\nRESET = Style.RESET_ALL\nSYSTEM_COLOR = Fore.CYAN\nUSER_COLOR = Fore.GREEN\nASSISTANT_COLOR = Fore.YELLOW\nPIPE_COLOR = Fore.BLUE\n\n# Set up logging\n# logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\nlogger = logging.getLogger(__name__)\n\nhas_logged_version_statement = False\n\n_version_check_lock = threading.Lock()\n_has_logged_version_statement = False",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 242,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 37,
            "community": null
        },
        "node_id": "util\\verbosity.py::1"
    },
    {
        "id": "util\\verbosity.py::2",
        "input_type": "file",
        "content": "def check_version_and_log():\n    global _has_logged_version_statement\n    if _version_check_lock.acquire(blocking=False):\n        try:\n            if not _has_logged_version_statement:\n\n                import ell\n\n                try:\n                    response = requests.get(\"https://version.ell.so/ell-ai/pypi\", timeout=0.15)\n                    if response.status_code == 200:\n                        latest_version = response.text.strip()\n                        if latest_version != ell.__version__:\n                            print(f\"{Fore.YELLOW}\u2554\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2557\")\n                            print(f\"{Fore.YELLOW}\u2551 {Fore.GREEN}A new version of ELL is available: {Fore.CYAN}{latest_version:<29}{Fore.YELLOW}\u2551\")\n                            print(f\"{Fore.YELLOW}\u2551 {Fore.GREEN}You can update by running:{Fore.YELLOW}                                      \u2551\")\n                            print(f\"{Fore.YELLOW}\u2551 {Fore.CYAN}pip install --upgrade ell-ai{Fore.YELLOW}                                    \u2551\")\n                            print(f\"{Fore.YELLOW}\u255a\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u2550\u255d{Style.RESET_ALL}\")\n                except requests.RequestException:\n                    pass  # Silently handle any network-related errors\n                _has_logged_version_statement = True\n        finally:\n            _version_check_lock.release()",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 316,
            "span_ids": [
                "check_version_and_log"
            ],
            "start_line": 39,
            "end_line": 61,
            "community": null
        },
        "node_id": "util\\verbosity.py::2"
    },
    {
        "id": "util\\verbosity.py::3",
        "input_type": "file",
        "content": "@lru_cache(maxsize=128)\ndef compute_color(invoking_lmp: LMP) -> str:\n    \"\"\"Compute and cache a consistent color for a given LMP.\"\"\"\n    name_hash = hashlib.md5(invoking_lmp.__name__.encode()).hexdigest()\n    color_index = int(name_hash, 16) % len(ELL_COLORS)\n    return list(ELL_COLORS.values())[color_index]\n\ndef format_arg(arg: Any, max_length: int = 8) -> str:\n    \"\"\"Format an argument for display with customizable max length.\"\"\"\n    str_arg = str(arg)\n    return f\"{Fore.MAGENTA}{str_arg[:max_length]}..{Style.RESET_ALL}\" if len(str_arg) > max_length else f\"{Fore.MAGENTA}{str_arg}{Style.RESET_ALL}\"\n\ndef format_kwarg(key: str, value: Any, max_length: int = 8) -> str:\n    \"\"\"Format a keyword argument for display with customizable max length.\"\"\"\n    return f\"{Style.DIM}{key}{Style.RESET_ALL}={Fore.MAGENTA}{str(value)[:max_length]}..{Style.RESET_ALL}\"\n\ndef get_terminal_width() -> int:\n    \"\"\"Get the terminal width, defaulting to 80 if it can't be determined.\"\"\"\n    try:\n        return shutil.get_terminal_size((80, 20)).columns\n    except Exception:\n        logger.warning(\"Unable to determine terminal size. Defaulting to 80 columns.\")\n        return 80",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 301,
            "span_ids": [
                "compute_color",
                "format_kwarg",
                "format_arg",
                "get_terminal_width"
            ],
            "start_line": 64,
            "end_line": 86,
            "community": null
        },
        "node_id": "util\\verbosity.py::3"
    },
    {
        "id": "util\\verbosity.py::4",
        "input_type": "file",
        "content": "def wrap_text_with_prefix(message, width: int, prefix: str, subsequent_prefix: str, text_color: str) -> List[str]:\n    \"\"\"Wrap text while preserving the prefix and color for each line.\"\"\"\n    result = []\n    for i, content in enumerate(message.content):\n        wrapped_lines = []\n        if content.image and content.image.image:\n            wrapped_lines = plot_ascii(content.image.image, min(80, width - len(prefix)))\n        else:\n            if content.tool_result:\n                contnets_to_wrap = [ContentBlock(text=f\"ToolResult(tool_call_id={content.tool_result.tool_call_id}):\"), *content.tool_result.result]\n            else:\n                contnets_to_wrap = [content]\n\n            wrapped_lines = []\n            for c in contnets_to_wrap:\n                if c.image and c.image.image:\n                    block_wrapped_lines = plot_ascii(c.image.image, min(80, width - len(prefix)))\n                else:\n                    text = _content_to_text([c])\n                    paragraphs = text.split('\\n')\n                    wrapped_paragraphs = [textwrap.wrap(p, width=width - len(prefix)) for p in paragraphs]\n                    block_wrapped_lines = [line for paragraph in wrapped_paragraphs for line in paragraph]\n                wrapped_lines.extend(block_wrapped_lines)\n        if i == 0:\n            if wrapped_lines:\n                result.append(f\"{prefix}{text_color}{wrapped_lines[0]}{RESET}\")\n            else:\n                result.append(f\"{prefix}{text_color}{RESET}\")\n        else:\n            result.append(f\"{subsequent_prefix}{text_color}{wrapped_lines[0]}{RESET}\")\n        result.extend([f\"{subsequent_prefix}{text_color}{line}{RESET}\" for line in wrapped_lines[1:]])\n    return result",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 356,
            "span_ids": [
                "wrap_text_with_prefix"
            ],
            "start_line": 88,
            "end_line": 119,
            "community": null
        },
        "node_id": "util\\verbosity.py::4"
    },
    {
        "id": "util\\verbosity.py::5",
        "input_type": "file",
        "content": "def print_wrapped_messages(messages: List[Message], max_role_length: int, color: str, wrap_width: Optional[int] = None):\n    \"\"\"Print wrapped messages with proper indentation, customizable wrap width, and consistent ASCII piping.\"\"\"\n    terminal_width = get_terminal_width()\n    prefix = f\"{PIPE_COLOR}\u2502   \"\n    role_prefix = ' ' * (max_role_length + 2)\n    subsequent_prefix = f\"{PIPE_COLOR}\u2502   {role_prefix}\"\n    wrapping_width = wrap_width or (terminal_width - len(prefix))\n\n    for i, message in enumerate(messages):\n        role = message.role\n        m = message.content[0]\n\n\n        role_color = SYSTEM_COLOR if role == \"system\" else USER_COLOR if role == \"user\" else ASSISTANT_COLOR\n\n        role_line = f\"{prefix}{role_color}{role.rjust(max_role_length)}: {RESET}\"\n        wrapped_lines = wrap_text_with_prefix(message, wrapping_width - len(role_prefix), '', subsequent_prefix, role_color)\n\n        print(f\"{role_line}{wrapped_lines[0]}\")\n        for line in wrapped_lines[1:]:\n            print(line)\n\n        if i < len(messages) - 1:\n            print(f\"{PIPE_COLOR}\u2502{RESET}\")",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 255,
            "span_ids": [
                "print_wrapped_messages"
            ],
            "start_line": 121,
            "end_line": 144,
            "community": null
        },
        "node_id": "util\\verbosity.py::5"
    },
    {
        "id": "util\\verbosity.py::6",
        "input_type": "file",
        "content": "def model_usage_logger_pre(\n    invoking_lmp: LMP,\n    lmp_args: Tuple,\n    lmp_kwargs: Dict,\n    lmp_hash: str,\n    messages: List[Message],\n    arg_max_length: int = 8\n):\n    \"\"\"Log model usage before execution with customizable argument display length and ASCII box.\"\"\"\n    color =  compute_color(invoking_lmp)\n    formatted_args = [format_arg(arg, arg_max_length) for arg in lmp_args]\n    formatted_kwargs = [format_kwarg(key, lmp_kwargs[key], arg_max_length) for key in lmp_kwargs]\n    formatted_params = ', '.join(formatted_args + formatted_kwargs)\n\n    check_version_and_log()\n\n    terminal_width = get_terminal_width()\n\n    logger.info(f\"Invoking LMP: {invoking_lmp.__name__} (hash: {lmp_hash[:8]})\")\n\n    print(f\"{PIPE_COLOR}\u2554{'\u2550' * (terminal_width - 2)}\u2557{RESET}\")\n    print(f\"{PIPE_COLOR}\u2551 {color}{BOLD}{UNDERLINE}{invoking_lmp.__name__}{RESET}{color}({formatted_params}){RESET}\")\n    print(f\"{PIPE_COLOR}\u2560{'\u2550' * (terminal_width - 2)}\u2563{RESET}\")\n    print(f\"{PIPE_COLOR}\u2551 {BOLD}Prompt:{RESET}\")\n    print(f\"{PIPE_COLOR}\u255f{'\u2500' * (terminal_width - 2)}\u2562{RESET}\")\n\n    max_role_length = max(len(\"assistant\"), max(len(message.role) for message in messages))\n    print_wrapped_messages(messages, max_role_length, color)",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 344,
            "span_ids": [
                "model_usage_logger_pre"
            ],
            "start_line": 147,
            "end_line": 174,
            "community": null
        },
        "node_id": "util\\verbosity.py::6"
    },
    {
        "id": "util\\verbosity.py::7",
        "input_type": "file",
        "content": "def model_usage_logger_post_start(color: str = \"\", n: int = 1):\n    \"\"\"Log the start of model output with ASCII box.\"\"\"\n    terminal_width = get_terminal_width()\n    print(f\"{PIPE_COLOR}\u255f{'\u2500' * (terminal_width - 2)}\u2562{RESET}\")\n    print(f\"{PIPE_COLOR}\u2551 {BOLD}Output{f'[0 of {n}]' if n > 1 else ''}:{RESET}\")\n    print(f\"{PIPE_COLOR}\u255f{'\u2500' * (terminal_width - 2)}\u2562{RESET}\")\n    print(f\"{PIPE_COLOR}\u2502   {ASSISTANT_COLOR}assistant: {RESET}\", end='')\n    sys.stdout.flush()\n\nfrom contextlib import contextmanager",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 153,
            "span_ids": [
                "impl:26",
                "model_usage_logger_post_start"
            ],
            "start_line": 176,
            "end_line": 185,
            "community": null
        },
        "node_id": "util\\verbosity.py::7"
    },
    {
        "id": "util\\verbosity.py::8",
        "input_type": "file",
        "content": "@contextmanager\ndef model_usage_logger_post_intermediate( n: int = 1):\n    \"\"\"Context manager to log intermediate model output without wrapping, only indenting if necessary.\"\"\"\n    terminal_width = get_terminal_width()\n    prefix = f\"{PIPE_COLOR}\u2502   \"\n    subsequent_prefix = f\"{PIPE_COLOR}\u2502   {' ' * (len('assistant: '))}\"\n    chars_printed = len(subsequent_prefix)\n\n    def log_stream_chunk(stream_chunk: str , is_refusal: bool = False):\n        nonlocal chars_printed\n        if stream_chunk:\n            lines = stream_chunk.split('\\n')\n            for i, line in enumerate(lines):\n                if chars_printed + len(line) > terminal_width - 6:\n                    print()\n                    if i == 0:\n                        print(subsequent_prefix, end='')\n                        chars_printed = len(prefix)\n                    else:\n                        print(subsequent_prefix, end='')\n                        chars_printed = len(subsequent_prefix)\n                    print(line.lstrip(), end='')\n                else:\n                    print(line, end='')\n                chars_printed += len(line)\n\n                if i < len(lines) - 1:\n                    print()\n                    print(subsequent_prefix, end='')\n                    chars_printed = len(subsequent_prefix)  # Reset for new line\n            sys.stdout.flush()\n\n    try:\n        yield log_stream_chunk\n    finally:\n        pass",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 282,
            "span_ids": [
                "model_usage_logger_post_intermediate"
            ],
            "start_line": 187,
            "end_line": 222,
            "community": null
        },
        "node_id": "util\\verbosity.py::8"
    },
    {
        "id": "util\\verbosity.py::9",
        "input_type": "file",
        "content": "def model_usage_logger_post_end():\n    \"\"\"Log the end of model output with ASCII box closure.\"\"\"\n    terminal_width = get_terminal_width()\n    print(f\"\\n{PIPE_COLOR}\u255a{'\u2550' * (terminal_width - 2)}\u255d{RESET}\")\n\ndef set_log_level(level: str):\n    \"\"\"Set the logging level.\"\"\"\n    numeric_level = getattr(logging, level.upper(), None)\n    if not isinstance(numeric_level, int):\n        raise ValueError(f'Invalid log level: {level}')\n    logger.setLevel(numeric_level)",
        "filepath": "src\\ell\\util\\verbosity.py",
        "metadata": {
            "file_path": "src\\ell\\util\\verbosity.py",
            "file_name": "verbosity.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 110,
            "span_ids": [
                "set_log_level",
                "model_usage_logger_post_end"
            ],
            "start_line": 223,
            "end_line": 233,
            "community": null
        },
        "node_id": "util\\verbosity.py::9"
    },
    {
        "id": "openai_realtime\\__init__.py::1",
        "input_type": "file",
        "content": "from .client import RealtimeClient\nfrom .api import RealtimeAPI\nfrom .conversation import RealtimeConversation\nfrom .event_handler import RealtimeEventHandler\nfrom .utils import RealtimeUtils\n\n__all__ = [\n    \"RealtimeClient\",\n    \"RealtimeAPI\",\n    \"RealtimeConversation\",\n    \"RealtimeEventHandler\",\n    \"RealtimeUtils\"\n]",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\__init__.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\__init__.py",
            "file_name": "__init__.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 77,
            "span_ids": [
                "imports"
            ],
            "start_line": 1,
            "end_line": 14,
            "community": null
        },
        "node_id": "openai_realtime\\__init__.py::1"
    },
    {
        "id": "openai_realtime\\api.py::1",
        "input_type": "file",
        "content": "import asyncio\nimport json\nimport websockets\nfrom .event_handler import RealtimeEventHandler\nfrom .utils import RealtimeUtils\n\nclass RealtimeAPI(RealtimeEventHandler):\n    def __init__(self, url=None, api_key=None, dangerously_allow_api_key_in_browser=False, debug=False):\n        super().__init__()\n        self.default_url = 'wss://api.openai.com/v1/realtime'\n        self.url = url or self.default_url\n        self.api_key = api_key\n        self.debug = debug\n        self.ws = None\n\n    def is_connected(self):\n        return self.ws is not None and self.ws.open\n\n    def log(self, *args):\n        if self.debug:\n            print(*args)\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\api.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\api.py",
            "file_name": "api.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 154,
            "span_ids": [
                "RealtimeAPI.__init__",
                "RealtimeAPI.log",
                "RealtimeAPI.is_connected",
                "imports",
                "RealtimeAPI"
            ],
            "start_line": 1,
            "end_line": 22,
            "community": null
        },
        "node_id": "openai_realtime\\api.py::1"
    },
    {
        "id": "openai_realtime\\api.py::2",
        "input_type": "file",
        "content": "class RealtimeAPI(RealtimeEventHandler):\n\n    async def connect(self, model='gpt-4o-realtime-preview-2024-10-01'):\n        if self.is_connected():\n            raise Exception(\"Already connected\")\n\n        headers = {\n            'Authorization': f'Bearer {self.api_key}',\n            'OpenAI-Beta': 'realtime=v1'\n        }\n\n        self.ws = await websockets.connect(f\"{self.url}?model={model}\", extra_headers=headers)\n\n        self.log(f\"Connected to {self.url}\")\n\n        asyncio.create_task(self._message_handler())\n\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\api.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\api.py",
            "file_name": "api.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 122,
            "span_ids": [
                "RealtimeAPI.connect"
            ],
            "start_line": 24,
            "end_line": 39,
            "community": null
        },
        "node_id": "openai_realtime\\api.py::2"
    },
    {
        "id": "openai_realtime\\api.py::3",
        "input_type": "file",
        "content": "class RealtimeAPI(RealtimeEventHandler):\n\n    async def _message_handler(self):\n        try:\n            async for message in self.ws:\n                data = json.loads(message)\n                self.receive(data['type'], data)\n        except websockets.exceptions.ConnectionClosed:\n            self.disconnect()\n            self.dispatch('close', {'error': True})\n\n    def disconnect(self):\n        if self.ws:\n            asyncio.create_task(self.ws.close())\n            self.ws = None\n        return True\n\n    def receive(self, event_name, event):\n        self.log(\"received:\", event_name, event)\n        self.dispatch(f\"server.{event_name}\", event)\n        self.dispatch(\"server.*\", event)\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\api.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\api.py",
            "file_name": "api.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 140,
            "span_ids": [
                "RealtimeAPI._message_handler",
                "RealtimeAPI.disconnect",
                "RealtimeAPI.receive"
            ],
            "start_line": 41,
            "end_line": 60,
            "community": null
        },
        "node_id": "openai_realtime\\api.py::3"
    },
    {
        "id": "openai_realtime\\api.py::4",
        "input_type": "file",
        "content": "class RealtimeAPI(RealtimeEventHandler):\n\n    def send(self, event_name, data=None):\n        if not self.is_connected():\n            raise Exception(\"RealtimeAPI is not connected\")\n\n        data = data or {}\n        if not isinstance(data, dict):\n            raise ValueError(\"data must be a dictionary\")\n\n        event = {\n            \"event_id\": RealtimeUtils.generate_id(\"evt_\"),\n            \"type\": event_name,\n            **data\n        }\n\n        self.dispatch(f\"client.{event_name}\", event)\n        self.dispatch(\"client.*\", event)\n        self.log(\"sent:\", event_name, event)\n\n        asyncio.create_task(self.ws.send(json.dumps(event, ensure_ascii=False)))\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\api.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\api.py",
            "file_name": "api.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 143,
            "span_ids": [
                "RealtimeAPI.send"
            ],
            "start_line": 62,
            "end_line": 81,
            "community": null
        },
        "node_id": "openai_realtime\\api.py::4"
    },
    {
        "id": "openai_realtime\\client.py::1",
        "input_type": "file",
        "content": "import asyncio\nimport numpy as np\nfrom .event_handler import RealtimeEventHandler\nfrom .api import RealtimeAPI\nfrom .conversation import RealtimeConversation\nfrom .utils import RealtimeUtils\nimport json\n\nclass RealtimeClient(RealtimeEventHandler):\n    def __init__(self, url=None, api_key=None, instructions='', dangerously_allow_api_key_in_browser=False, debug=False):\n        super().__init__()\n        self.default_session_config = {\n            'modalities': ['text', 'audio'],\n            'instructions': instructions,\n            'voice': 'alloy',\n            'input_audio_format': 'pcm16',\n            'output_audio_format': 'pcm16',\n            'input_audio_transcription': None,\n            'turn_detection': None,\n            'tools': [],\n            'tool_choice': 'auto',\n            'temperature': 0.8,\n            'max_response_output_tokens': 4096,\n        }\n        self.session_config = {}\n        self.transcription_models = [{'model': 'whisper-1'}]\n        self.default_server_vad_config = {\n            'type': 'server_vad',\n            'threshold': 0.5,\n            'prefix_padding_ms': 300,\n            'silence_duration_ms': 200,\n        }\n        self.realtime = RealtimeAPI(url, api_key, dangerously_allow_api_key_in_browser, debug)\n        self.conversation = RealtimeConversation()\n        self._reset_config()\n        self._add_api_event_handlers()\n\n    def _reset_config(self):\n        self.session_created = False\n        self.tools = {}\n        self.session_config = self.default_session_config.copy()\n        self.input_audio_buffer = np.array([], dtype=np.int16)\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 350,
            "span_ids": [
                "RealtimeClient._reset_config",
                "RealtimeClient",
                "RealtimeClient.__init__",
                "imports"
            ],
            "start_line": 1,
            "end_line": 43,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::1"
    },
    {
        "id": "openai_realtime\\client.py::2",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n    def _add_api_event_handlers(self):\n        self.realtime.on('client.*', lambda event: self.dispatch('realtime.event', {\n            'time': RealtimeUtils.generate_id('time_'),\n            'source': 'client',\n            'event': event\n        }))\n        self.realtime.on('server.*', lambda event: self.dispatch('realtime.event', {\n            'time': RealtimeUtils.generate_id('time_'),\n            'source': 'server',\n            'event': event\n        }))\n        self.realtime.on('server.session.created', lambda _: setattr(self, 'session_created', True))\n\n        def handle_conversation_event(event, *args):\n            result = self.conversation.process_event(event, *args)\n            if result['item']:\n                self.dispatch('conversation.updated', result)\n            return result\n\n        self.realtime.on('server.response.created', handle_conversation_event)\n        self.realtime.on('server.response.output_item.added', handle_conversation_event)\n        self.realtime.on('server.response.content_part.added', handle_conversation_event)\n        self.realtime.on('server.input_audio_buffer.speech_started', lambda event: (\n            handle_conversation_event(event),\n            self.dispatch('conversation.interrupted', event)\n        ))\n        self.realtime.on('server.input_audio_buffer.speech_stopped', lambda event: \n            handle_conversation_event(event, self.input_audio_buffer)\n        )\n        self.realtime.on('server.conversation.item.created', lambda event: (\n            handle_conversation_event(event),\n            self.dispatch('conversation.item.appended', {'item': event['item']})\n        ))\n        self.realtime.on('server.conversation.item.truncated', handle_conversation_event)\n        self.realtime.on('server.conversation.item.deleted', handle_conversation_event)\n        self.realtime.on('server.conversation.item.input_audio_transcription.completed', handle_conversation_event)\n        self.realtime.on('server.response.audio_transcript.delta', handle_conversation_event)\n        self.realtime.on('server.response.audio.delta', handle_conversation_event)\n        self.realtime.on('server.response.text.delta', handle_conversation_event)\n        self.realtime.on('server.response.function_call_arguments.delta', handle_conversation_event)\n        def handle_output_item_done( event):\n            handle_conversation_event(event)\n            item = event.get('item', {})\n\n            if item.get('status') == 'completed':\n                self.dispatch('conversation.item.completed', {'item': item})\n\n            formatted = item.get('formatted', {})\n            tool = formatted.get('tool') if isinstance(formatted, dict) else None\n\n            if tool:\n                asyncio.create_task(self._call_tool(tool))\n        self.realtime.on('server.response.output_item.done', handle_output_item_done)",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 573,
            "span_ids": [
                "RealtimeClient._add_api_event_handlers"
            ],
            "start_line": 45,
            "end_line": 97,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::2"
    },
    {
        "id": "openai_realtime\\client.py::3",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n\n\n    def is_connected(self):\n        return self.realtime.is_connected() and self.session_created\n\n    def reset(self):\n        self.disconnect()\n        self.clear_event_handlers()\n        self.realtime.clear_event_handlers()\n        self._reset_config()\n        self._add_api_event_handlers()\n        return True\n\n    async def connect(self):\n        if self.is_connected():\n            raise Exception(\"Already connected, use .disconnect() first\")\n        await self.realtime.connect()\n        self.update_session()\n        return True\n\n    async def wait_for_session_created(self):\n        if not self.realtime.is_connected():\n            raise Exception(\"Not connected, use .connect() first\")\n        while not self.session_created:\n            await asyncio.sleep(0.001)\n        return True\n\n    def disconnect(self):\n        self.session_created = False\n        self.conversation.clear()\n        if self.realtime.is_connected():\n            self.realtime.disconnect()\n\n    def get_turn_detection_type(self):\n        turn_detection = self.session_config.get('turn_detection')\n        if isinstance(turn_detection, dict):\n            return turn_detection.get('type')\n        return None",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 233,
            "span_ids": [
                "RealtimeClient.wait_for_session_created",
                "RealtimeClient.connect",
                "RealtimeClient.get_turn_detection_type",
                "RealtimeClient.is_connected",
                "RealtimeClient.reset",
                "RealtimeClient.disconnect"
            ],
            "start_line": 101,
            "end_line": 136,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::3"
    },
    {
        "id": "openai_realtime\\client.py::4",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n    def add_tool(self, definition, handler):\n        if not definition.get('name'):\n            raise ValueError(\"Missing tool name in definition\")\n        name = definition['name']\n        if name in self.tools:\n            raise ValueError(f\"Tool '{name}' already added. Please use .remove_tool('{name}') before trying to add again.\")\n        if not callable(handler):\n            raise ValueError(f\"Tool '{name}' handler must be a function\")\n        self.tools[name] = {'definition': definition, 'handler': handler}\n        self.update_session()\n        return self.tools[name]",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 125,
            "span_ids": [
                "RealtimeClient.add_tool"
            ],
            "start_line": 138,
            "end_line": 148,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::4"
    },
    {
        "id": "openai_realtime\\client.py::5",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n    def remove_tool(self, name):\n        if name not in self.tools:\n            raise ValueError(f\"Tool '{name}' does not exist, cannot be removed.\")\n        del self.tools[name]\n        return True\n\n    def delete_item(self, id):\n        self.realtime.send('conversation.item.delete', {'item_id': id})\n        return True\n\n    def update_session(self, **kwargs):\n        self.session_config.update(kwargs)\n        use_tools = [\n            {**tool.get('definition', {}), 'type': 'function'}\n            for tool in self.tools.values()\n        ]\n        session = {**self.session_config, 'tools': use_tools}\n        if self.realtime.is_connected():\n            self.realtime.send('session.update', {'session': session})\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 168,
            "span_ids": [
                "RealtimeClient.delete_item",
                "RealtimeClient.remove_tool",
                "RealtimeClient.update_session"
            ],
            "start_line": 150,
            "end_line": 169,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::5"
    },
    {
        "id": "openai_realtime\\client.py::6",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n    def send_user_message_content(self, content=None):\n        content = content or []\n        for c in content:\n            if c['type'] == 'input_audio':\n                if isinstance(c['audio'], (np.ndarray, bytes)):\n                    c['audio'] = RealtimeUtils.array_buffer_to_base64(c['audio'])\n        if content:\n            self.realtime.send('conversation.item.create', {\n                'item': {\n                    'type': 'message',\n                    'role': 'user',\n                    'content': content\n                }\n            })\n        self.create_response()\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 126,
            "span_ids": [
                "RealtimeClient.send_user_message_content"
            ],
            "start_line": 171,
            "end_line": 186,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::6"
    },
    {
        "id": "openai_realtime\\client.py::7",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n    def append_input_audio(self, array_buffer):\n        if len(array_buffer) > 0:\n            self.realtime.send('input_audio_buffer.append', {\n                'audio': RealtimeUtils.array_buffer_to_base64(array_buffer)\n            })\n            self.input_audio_buffer = RealtimeUtils.merge_int16_arrays(\n                self.input_audio_buffer,\n                array_buffer\n            )\n        return True\n\n    def create_response(self):\n        if self.get_turn_detection_type() is None and len(self.input_audio_buffer) > 0:\n            self.realtime.send('input_audio_buffer.commit')\n            self.conversation.queue_input_audio(self.input_audio_buffer)\n            self.input_audio_buffer = np.array([], dtype=np.int16)\n        self.realtime.send('response.create')\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 164,
            "span_ids": [
                "RealtimeClient.create_response",
                "RealtimeClient.append_input_audio"
            ],
            "start_line": 188,
            "end_line": 205,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::7"
    },
    {
        "id": "openai_realtime\\client.py::8",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n    def cancel_response(self, id=None, sample_count=0):\n        if not id:\n            self.realtime.send('response.cancel')\n            return {'item': None}\n        item = self.conversation.get_item(id)\n        if not item:\n            raise ValueError(f\"Could not find item '{id}'\")\n        if item['type'] != 'message' or item['role'] != 'assistant':\n            raise ValueError(\"Can only cancel response messages with type 'message' and role 'assistant'\")\n        self.realtime.send('response.cancel')\n        audio_index = next((i for i, c in enumerate(item['content']) if c['type'] == 'audio'), -1)\n        if audio_index == -1:\n            raise ValueError(\"Could not find audio on item to cancel\")\n        self.realtime.send('conversation.item.truncate', {\n            'item_id': id,\n            'content_index': audio_index,\n            'audio_end_ms': int((sample_count / self.conversation.default_frequency) * 1000)\n        })\n        return {'item': item}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 227,
            "span_ids": [
                "RealtimeClient.cancel_response"
            ],
            "start_line": 207,
            "end_line": 225,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::8"
    },
    {
        "id": "openai_realtime\\client.py::9",
        "input_type": "file",
        "content": "class RealtimeClient(RealtimeEventHandler):\n\n    async def wait_for_next_item(self):\n        event = await self.wait_for_next('conversation.item.appended')\n        return {'item': event['item']}\n\n    async def wait_for_next_completed_item(self):\n        event = await self.wait_for_next('conversation.item.completed')\n        return {'item': event['item']}\n\n    async def _call_tool(self, tool):\n        try:\n            json_arguments = json.loads(tool['arguments'])\n            tool_config = self.tools.get(tool['name'])\n            if not tool_config:\n                raise ValueError(f\"Tool '{tool['name']}' has not been added\")\n            result = await tool_config['handler'](json_arguments)\n            self.realtime.send('conversation.item.create', {\n                'item': {\n                    'type': 'function_call_output',\n                    'call_id': tool['call_id'],\n                    'output': json.dumps(result, ensure_ascii=False)\n                }\n            })\n        except Exception as e:\n            self.realtime.send('conversation.item.create', {\n                'item': {\n                    'type': 'function_call_output',\n                    'call_id': tool['call_id'],\n                    'output': json.dumps({'error': str(e)}, ensure_ascii=False)\n                }\n            })\n        self.create_response()",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\client.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\client.py",
            "file_name": "client.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 260,
            "span_ids": [
                "RealtimeClient.wait_for_next_completed_item",
                "RealtimeClient._call_tool",
                "RealtimeClient.wait_for_next_item"
            ],
            "start_line": 227,
            "end_line": 257,
            "community": null
        },
        "node_id": "openai_realtime\\client.py::9"
    },
    {
        "id": "openai_realtime\\conversation.py::1",
        "input_type": "file",
        "content": "import numpy as np\nimport json\nfrom .utils import RealtimeUtils\nimport copy\n\nclass RealtimeConversation:\n    def __init__(self):\n        self.default_frequency = 24000  # 24,000 Hz\n        self.clear()\n\n    def clear(self):\n        self.item_lookup = {}\n        self.items = []\n        self.response_lookup = {}\n        self.responses = []\n        self.queued_speech_items = {}\n        self.queued_transcript_items = {}\n        self.queued_input_audio = None\n        return True\n\n    def queue_input_audio(self, input_audio):\n        self.queued_input_audio = input_audio\n        return input_audio\n\n    def process_event(self, event, *args):\n        if 'event_id' not in event:\n            raise ValueError(\"Missing 'event_id' on event\")\n        if 'type' not in event:\n            raise ValueError(\"Missing 'type' on event\")\n\n        event_processor = getattr(self, f\"_process_{event['type'].replace('.', '_')}\", None)\n        if not event_processor:\n            raise ValueError(f\"Missing conversation event processor for '{event['type']}'\")\n\n        return event_processor(event, *args)\n\n    def get_item(self, id):\n        return self.item_lookup.get(id)\n\n    def get_items(self):\n        return self.items.copy()",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 268,
            "span_ids": [
                "RealtimeConversation.__init__",
                "RealtimeConversation.get_items",
                "RealtimeConversation.process_event",
                "RealtimeConversation.get_item",
                "RealtimeConversation",
                "RealtimeConversation.clear",
                "imports",
                "RealtimeConversation.queue_input_audio"
            ],
            "start_line": 1,
            "end_line": 41,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::1"
    },
    {
        "id": "openai_realtime\\conversation.py::2",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_conversation_item_created(self, event):\n        item = event['item']\n        new_item = copy.deepcopy(item)\n        if new_item['id'] not in self.item_lookup:\n            self.item_lookup[new_item['id']] = new_item\n            self.items.append(new_item)\n\n        new_item['formatted'] = {\n            'audio': np.array([], dtype=np.int16),\n            'text': '',\n            'transcript': ''\n        }\n\n        if new_item['type'] == 'message':\n            if new_item['role'] == 'user':\n                new_item['status'] = 'completed'\n                if self.queued_input_audio is not None:\n                    new_item['formatted']['audio'] = self.queued_input_audio\n                    self.queued_input_audio = None\n            else:\n                new_item['status'] = 'in_progress'\n        elif new_item['type'] == 'function_call':\n            new_item['formatted']['tool'] = {\n                'type': 'function',\n                'name': new_item['name'],\n                'call_id': new_item['call_id'],\n                'arguments': ''\n            }\n            new_item['status'] = 'in_progress'\n        elif new_item['type'] == 'function_call_output':\n            new_item['status'] = 'completed'\n            new_item['formatted']['output'] = new_item['output']\n\n        if new_item.get('content'):\n            text_content = [c for c in new_item['content'] if c['type'] in ['text', 'input_text']]\n            for content in text_content:\n                new_item['formatted']['text'] += content['text']\n\n        if new_item['id'] in self.queued_speech_items:\n            new_item['formatted']['audio'] = self.queued_speech_items[new_item['id']]['audio']\n            del self.queued_speech_items[new_item['id']]\n\n        if new_item['id'] in self.queued_transcript_items:\n            new_item['formatted']['transcript'] = self.queued_transcript_items[new_item['id']]['transcript']\n            del self.queued_transcript_items[new_item['id']]\n\n        return {'item': new_item, 'delta': None}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 455,
            "span_ids": [
                "RealtimeConversation._process_conversation_item_created"
            ],
            "start_line": 43,
            "end_line": 89,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::2"
    },
    {
        "id": "openai_realtime\\conversation.py::3",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_conversation_item_truncated(self, event):\n        item_id, audio_end_ms = event['item_id'], event['audio_end_ms']\n        item = self.item_lookup.get(item_id)\n        if not item:\n            raise ValueError(f\"item.truncated: Item '{item_id}' not found\")\n\n        end_index = int((audio_end_ms * self.default_frequency) / 1000)\n        item['formatted']['transcript'] = ''\n        item['formatted']['audio'] = item['formatted']['audio'][:end_index]\n        return {'item': item, 'delta': None}\n\n    def _process_conversation_item_deleted(self, event):\n        item_id = event['item_id']\n        item = self.item_lookup.get(item_id)\n        if not item:\n            raise ValueError(f\"item.deleted: Item '{item_id}' not found\")\n\n        del self.item_lookup[item['id']]\n        self.items = [i for i in self.items if i['id'] != item['id']]\n        return {'item': item, 'delta': None}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 222,
            "span_ids": [
                "RealtimeConversation._process_conversation_item_deleted",
                "RealtimeConversation._process_conversation_item_truncated"
            ],
            "start_line": 91,
            "end_line": 110,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::3"
    },
    {
        "id": "openai_realtime\\conversation.py::4",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_conversation_item_input_audio_transcription_completed(self, event):\n        item_id, content_index, transcript = event['item_id'], event['content_index'], event['transcript']\n        item = self.item_lookup.get(item_id)\n        formatted_transcript = transcript or ' '\n\n        if not item:\n            self.queued_transcript_items[item_id] = {'transcript': formatted_transcript}\n            return {'item': None, 'delta': None}\n\n        item['content'][content_index]['transcript'] = transcript\n        item['formatted']['transcript'] = formatted_transcript\n        return {'item': item, 'delta': {'transcript': transcript}}\n\n    def _process_input_audio_buffer_speech_started(self, event):\n        item_id, audio_start_ms = event['item_id'], event['audio_start_ms']\n        self.queued_speech_items[item_id] = {'audio_start_ms': audio_start_ms}\n        return {'item': None, 'delta': None}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 208,
            "span_ids": [
                "RealtimeConversation._process_input_audio_buffer_speech_started",
                "RealtimeConversation._process_conversation_item_input_audio_transcription_completed"
            ],
            "start_line": 112,
            "end_line": 128,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::4"
    },
    {
        "id": "openai_realtime\\conversation.py::5",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_input_audio_buffer_speech_stopped(self, event, input_audio_buffer):\n        item_id, audio_end_ms = event['item_id'], event['audio_end_ms']\n        speech = self.queued_speech_items[item_id]\n        speech['audio_end_ms'] = audio_end_ms\n        if input_audio_buffer is not None:\n            start_index = int((speech['audio_start_ms'] * self.default_frequency) / 1000)\n            end_index = int((speech['audio_end_ms'] * self.default_frequency) / 1000)\n            speech['audio'] = input_audio_buffer[start_index:end_index]\n        return {'item': None, 'delta': None}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 146,
            "span_ids": [
                "RealtimeConversation._process_input_audio_buffer_speech_stopped"
            ],
            "start_line": 130,
            "end_line": 138,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::5"
    },
    {
        "id": "openai_realtime\\conversation.py::6",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_response_created(self, event):\n        response = event['response']\n        if response['id'] not in self.response_lookup:\n            self.response_lookup[response['id']] = response\n            self.responses.append(response)\n        return {'item': None, 'delta': None}\n\n    def _process_response_output_item_added(self, event):\n        response_id, item = event['response_id'], event['item']\n        response = self.response_lookup.get(response_id)\n        if not response:\n            raise ValueError(f\"response.output_item.added: Response '{response_id}' not found\")\n        response['output'].append(item['id'])\n        return {'item': None, 'delta': None}\n\n    def _process_response_output_item_done(self, event):\n        item = event['item']\n        if not item:\n            raise ValueError(\"response.output_item.done: Missing 'item'\")\n        found_item = self.item_lookup.get(item['id'])\n        if not found_item:\n            raise ValueError(f\"response.output_item.done: Item '{item['id']}' not found\")\n        found_item['status'] = item['status']\n        return {'item': found_item, 'delta': None}\n\n    def _process_response_content_part_added(self, event):\n        item_id, part = event['item_id'], event['part']\n        item = self.item_lookup.get(item_id)\n        if not item:\n            raise ValueError(f\"response.content_part.added: Item '{item_id}' not found\")\n        item['content'].append(part)\n        return {'item': item, 'delta': None}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 326,
            "span_ids": [
                "RealtimeConversation._process_response_output_item_done",
                "RealtimeConversation._process_response_created",
                "RealtimeConversation._process_response_content_part_added",
                "RealtimeConversation._process_response_output_item_added"
            ],
            "start_line": 140,
            "end_line": 171,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::6"
    },
    {
        "id": "openai_realtime\\conversation.py::7",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_response_audio_transcript_delta(self, event):\n        item_id, content_index, delta = event['item_id'], event['content_index'], event['delta']\n        item = self.item_lookup.get(item_id)\n        if not item:\n            raise ValueError(f\"response.audio_transcript.delta: Item '{item_id}' not found\")\n        item['content'][content_index]['transcript'] += delta\n        item['formatted']['transcript'] += delta\n        return {'item': item, 'delta': {'transcript': delta}}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 116,
            "span_ids": [
                "RealtimeConversation._process_response_audio_transcript_delta"
            ],
            "start_line": 173,
            "end_line": 180,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::7"
    },
    {
        "id": "openai_realtime\\conversation.py::8",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_response_audio_delta(self, event):\n        item_id, content_index, delta = event['item_id'], event['content_index'], event['delta']\n        item = self.item_lookup.get(item_id)\n        if not item:\n            raise ValueError(f\"response.audio.delta: Item '{item_id}' not found\")\n        array_buffer = RealtimeUtils.base64_to_array_buffer(delta)\n        append_values = np.frombuffer(array_buffer, dtype=np.int16)\n        item['formatted']['audio'] = np.concatenate([item['formatted']['audio'], append_values])\n        return {'item': item, 'delta': {'audio': append_values}}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 136,
            "span_ids": [
                "RealtimeConversation._process_response_audio_delta"
            ],
            "start_line": 182,
            "end_line": 190,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::8"
    },
    {
        "id": "openai_realtime\\conversation.py::9",
        "input_type": "file",
        "content": "class RealtimeConversation:\n\n    def _process_response_text_delta(self, event):\n        item_id, content_index, delta = event['item_id'], event['content_index'], event['delta']\n        item = self.item_lookup.get(item_id)\n        if not item:\n            raise ValueError(f\"response.text.delta: Item '{item_id}' not found\")\n        item['content'][content_index]['text'] += delta\n        item['formatted']['text'] += delta\n        return {'item': item, 'delta': {'text': delta}}\n\n    def _process_response_function_call_arguments_delta(self, event):\n        item_id, delta = event['item_id'], event['delta']\n        item = self.item_lookup.get(item_id)\n        if not item:\n            raise ValueError(f\"response.function_call_arguments.delta: Item '{item_id}' not found\")\n        item['arguments'] += delta\n        item['formatted']['tool']['arguments'] += delta\n        return {'item': item, 'delta': {'arguments': delta}}",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\conversation.py",
            "file_name": "conversation.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 206,
            "span_ids": [
                "RealtimeConversation._process_response_function_call_arguments_delta",
                "RealtimeConversation._process_response_text_delta"
            ],
            "start_line": 192,
            "end_line": 208,
            "community": null
        },
        "node_id": "openai_realtime\\conversation.py::9"
    },
    {
        "id": "openai_realtime\\event_handler.py::1",
        "input_type": "file",
        "content": "import asyncio\nfrom typing import Callable, Dict, List, Any\n\nclass RealtimeEventHandler:\n    def __init__(self):\n        self.event_handlers: Dict[str, List[Callable]] = {}\n        self.next_event_handlers: Dict[str, List[Callable]] = {}\n\n    def clear_event_handlers(self):\n        self.event_handlers.clear()\n        self.next_event_handlers.clear()\n        return True\n\n    def on(self, event_name: str, callback: Callable = None):\n        def decorator(func):\n            if event_name not in self.event_handlers:\n                self.event_handlers[event_name] = []\n            self.event_handlers[event_name].append(func)\n            return func\n\n        if callback is None:\n            return decorator\n        else:\n            return decorator(callback)\n\n    def on_next(self, event_name: str, callback: Callable):\n        if event_name not in self.next_event_handlers:\n            self.next_event_handlers[event_name] = []\n        self.next_event_handlers[event_name].append(callback)\n\n    def off(self, event_name: str, callback: Callable = None):\n        if event_name in self.event_handlers:\n            if callback:\n                self.event_handlers[event_name].remove(callback)\n            else:\n                del self.event_handlers[event_name]\n        return True\n\n    def off_next(self, event_name: str, callback: Callable = None):\n        if event_name in self.next_event_handlers:\n            if callback:\n                self.next_event_handlers[event_name].remove(callback)\n            else:\n                del self.next_event_handlers[event_name]\n        return True\n\n    async def wait_for_next(self, event_name: str, timeout: float = None):\n        next_event = None\n        def set_next_event(event):\n            nonlocal next_event\n            next_event = event\n\n        self.on_next(event_name, set_next_event)\n\n        start_time = asyncio.get_event_loop().time()\n        while not next_event:\n            if timeout and asyncio.get_event_loop().time() - start_time > timeout:\n                return None\n            await asyncio.sleep(0.001)\n\n        return next_event\n\n    def dispatch(self, event_name: str, event: Any):\n        handlers = self.event_handlers.get(event_name, []).copy()\n        for handler in handlers:\n            handler(event)\n\n        next_handlers = self.next_event_handlers.pop(event_name, [])\n        for next_handler in next_handlers:\n            next_handler(event)\n\n        return True",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\event_handler.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\event_handler.py",
            "file_name": "event_handler.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 483,
            "span_ids": [
                "RealtimeEventHandler.off",
                "RealtimeEventHandler.off_next",
                "RealtimeEventHandler",
                "RealtimeEventHandler.clear_event_handlers",
                "RealtimeEventHandler.wait_for_next",
                "RealtimeEventHandler.on_next",
                "imports",
                "RealtimeEventHandler.on",
                "RealtimeEventHandler.__init__",
                "RealtimeEventHandler.dispatch"
            ],
            "start_line": 1,
            "end_line": 72,
            "community": null
        },
        "node_id": "openai_realtime\\event_handler.py::1"
    },
    {
        "id": "openai_realtime\\utils.py::1",
        "input_type": "file",
        "content": "import base64\nimport numpy as np\n\nclass RealtimeUtils:\n    @staticmethod\n    def float_to_16bit_pcm(float32_array):\n        int16_array = (np.clip(float32_array, -1, 1) * 32767).astype(np.int16)\n        return int16_array.tobytes()\n\n    @staticmethod\n    def base64_to_array_buffer(base64_string):\n        return base64.b64decode(base64_string)\n\n    @staticmethod\n    def array_buffer_to_base64(array_buffer):\n        if isinstance(array_buffer, np.ndarray):\n            if array_buffer.dtype == np.float32:\n                array_buffer = RealtimeUtils.float_to_16bit_pcm(array_buffer)\n            elif array_buffer.dtype == np.int16:\n                array_buffer = array_buffer.tobytes()\n        return base64.b64encode(array_buffer).decode('utf-8')\n\n    @staticmethod\n    def merge_int16_arrays(left, right):\n        if isinstance(left, bytes):\n            left = np.frombuffer(left, dtype=np.int16)\n        if isinstance(right, bytes):\n            right = np.frombuffer(right, dtype=np.int16)\n        if not isinstance(left, np.ndarray) or not isinstance(right, np.ndarray):\n            raise ValueError(\"Both items must be numpy arrays or bytes objects\")\n        return np.concatenate((left, right))\n\n    @staticmethod\n    def generate_id(prefix, length=21):\n        import random\n        chars = '123456789ABCDEFGHJKLMNPQRSTUVWXYZabcdefghijkmnopqrstuvwxyz'\n        return prefix + ''.join(random.choice(chars) for _ in range(length - len(prefix)))",
        "filepath": "x\\openai_realtime\\src\\openai_realtime\\utils.py",
        "metadata": {
            "file_path": "x\\openai_realtime\\src\\openai_realtime\\utils.py",
            "file_name": "utils.py",
            "file_type": "text/x-python",
            "category": "implementation",
            "tokens": 327,
            "span_ids": [
                "RealtimeUtils.base64_to_array_buffer",
                "RealtimeUtils.generate_id",
                "RealtimeUtils.array_buffer_to_base64",
                "imports",
                "RealtimeUtils.merge_int16_arrays",
                "RealtimeUtils",
                "RealtimeUtils.float_to_16bit_pcm"
            ],
            "start_line": 1,
            "end_line": 37,
            "community": null
        },
        "node_id": "openai_realtime\\utils.py::1"
    }
]