Send metadata with llm response

[Deepak-Kesavan]

What

• Send metadata with llm response
• Included cost calculation for LLM and embedding
• Added index for usage table on field run_id

Why

• To track token usage
• To track cost per run_id

How

• Send it as new key metadata along with LLM response
• Cost calculation is done after the callback from llama-index using litellm JSON.

Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)

• Mostly no. Because it is just addition of new key and moved existing data like run_id inside it.

Database Migrations

Migrations for 'usage': usage/migrations/0003_usage_cost_in_dollars_usage_llm_usage_reason_and_more.py

• Add field cost_in_dollars to usage
• Add field llm_usage_reason to usage
• Alter field usage_type on usage
• Create index token_usage_run_id_cd3578_idx on field(s) run_id of model usage

Env Config

Backend

STRUCTURE_TOOL_IMAGE_URL="docker:unstract/tool-structure:0.0.27"
STRUCTURE_TOOL_IMAGE_TAG="0.0.27"

Platform Service

MODEL_PRICES_URL="https://raw.githubusercontent.com/BerriAI/litellm/main/model_prices_and_context_window.json"
MODEL_PRICES_TTL_IN_DAYS=7
MODEL_PRICES_FILE_PATH="/tmp/model_prices.json"

Relevant Docs

-

Related Issues or PRs

• https://github.com/Zipstack/unstract-sdk/pull/61

Dependencies Versions

-

Notes on Testing

-

Screenshots

Checklist

I have read and understood the [Contribution Guidelines]().

[Deepak-Kesavan]

PR looks good to me. Except for 2 concerns

• Accessing DB from prompt_service. However, may not be a good idea to talk to platform service for this reason
• Seconding Chandru's comment on keeping usage_reason to be more generic field if possible.

Regarding keeping the usage_reason to be more generic field. It is done this way because we don't need separate token count of indexing and retrieval. Also this is like a temporary code made to quickly show the LLM usage split up. The long term solution is the logging implementation that Chandru is working on.

[sonarcloud[bot]]

Quality Gate passed

Issues
0 New issues
0 Accepted issues

Measures
0 Security Hotspots
0.0% Coverage on New Code
0.0% Duplication on New Code

See analysis details on SonarCloud

[github-actions[bot]]

filepath	function	$$\textcolor{#23d18b}{\tt{passed}}$$	SUBTOTAL
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_logs}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_cleanup}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_cleanup\_skip}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_client\_init}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_get\_image\_exists}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_get\_image}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_get\_container\_run\_config}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_get\_container\_run\_config\_without\_mount}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{worker/src/unstract/worker/clients/test\_docker.py}}$$	$$\textcolor{#23d18b}{\tt{test\_run\_container}}$$	$$\textcolor{#23d18b}{\tt{1}}$$	$$\textcolor{#23d18b}{\tt{1}}$$
$$\textcolor{#23d18b}{\tt{TOTAL}}$$		$$\textcolor{#23d18b}{\tt{9}}$$	$$\textcolor{#23d18b}{\tt{9}}$$