Send metadata with llm response (#413)

* Added metadata field which sends usage details with LLM response

* structure tool version bump

* Reverted prompt version manager change

* Minor refactor

* Updated logs

* Minor improvement

* Renamed function

* Added token cost calculation

* Updated SDK and adapters version

* Fixed security hotspot

* Tool version bump

* SDK and Adapters version bump

* SDK and Adapters version bump in backend

* Minor Fixes

Fixed missing module issues
Remove numpy usage and replaced with logic
Addressed review comments

backend/prompt_studio/prompt_studio_core/prompt_studio_helper.py

@@ -445,7 +445,7 @@ def prompt_responder(
                 )
 
                 OutputManagerHelper.handle_prompt_output_update(
-                    run_id=response[TSPKeys.RUN_ID],
+                    run_id=run_id,
                     prompts=prompts,
                     outputs=response["output"],
                     document_id=document_id,
@@ -521,9 +521,9 @@ def prompt_responder(
                 )
 
                 OutputManagerHelper.handle_prompt_output_update(
-                    run_id=response[TSPKeys.RUN_ID],
+                    run_id=run_id,
                     prompts=prompts,
-                    outputs=response[TSPKeys.SINGLE_PASS_EXTRACTION],
+                    outputs=response[TSPKeys.OUTPUT],
                     document_id=document_id,
                     is_single_pass_extract=True,
                 )

backend/pyproject.toml

@@ -31,8 +31,8 @@ dependencies = [
     "python-socketio==5.9.0", # For log_events
     "social-auth-app-django==5.3.0", # For OAuth
     "social-auth-core==4.4.2", # For OAuth
-    "unstract-sdk~=0.32.0",
-    "unstract-adapters~=0.19.0",
+    "unstract-sdk~=0.33.1",
+    "unstract-adapters~=0.19.2",
     # ! IMPORTANT!
     # Indirect local dependencies usually need to be added in their own projects
     # as: https://pdm-project.org/latest/usage/dependency/#local-dependencies.

backend/sample.env

@@ -89,9 +89,9 @@ PROMPT_PORT=3003
 PROMPT_STUDIO_FILE_PATH=/app/prompt-studio-data
 
 # Structure Tool
-STRUCTURE_TOOL_IMAGE_URL="docker:unstract/tool-structure:0.0.26"
+STRUCTURE_TOOL_IMAGE_URL="docker:unstract/tool-structure:0.0.27"
 STRUCTURE_TOOL_IMAGE_NAME="unstract/tool-structure"
-STRUCTURE_TOOL_IMAGE_TAG="0.0.26"
+STRUCTURE_TOOL_IMAGE_TAG="0.0.27"
 
 # Feature Flags
 EVALUATION_SERVER_IP=localhost

backend/usage/migrations/0003_usage_cost_in_dollars_usage_llm_usage_reason_and_more.py

@@ -0,0 +1,49 @@
+# Generated by Django 4.2.1 on 2024-06-27 08:08
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("usage", "0002_alter_usage_adapter_instance_id_and_more"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="usage",
+            name="cost_in_dollars",
+            field=models.FloatField(
+                db_comment="Total number of tokens used", default=0.0
+            ),
+            preserve_default=False,
+        ),
+        migrations.AddField(
+            model_name="usage",
+            name="llm_usage_reason",
+            field=models.CharField(
+                blank=True,
+                choices=[
+                    ("extraction", "Extraction"),
+                    ("challenge", "Challenge"),
+                    ("summarize", "Summarize"),
+                ],
+                db_comment="Reason for LLM usage. Empty if usage_type is 'embedding'. ",
+                max_length=255,
+                null=True,
+            ),
+        ),
+        migrations.AlterField(
+            model_name="usage",
+            name="usage_type",
+            field=models.CharField(
+                choices=[("llm", "LLM Usage"), ("embedding", "Embedding Usage")],
+                db_comment="Type of usage, either 'llm' or 'embedding'",
+                max_length=255,
+            ),
+        ),
+        migrations.AddIndex(
+            model_name="usage",
+            index=models.Index(fields=["run_id"], name="token_usage_run_id_cd3578_idx"),
+        ),
+    ]

backend/usage/models.py

@@ -4,6 +4,17 @@
 from utils.models.base_model import BaseModel
 
 
+class UsageType(models.TextChoices):
+    LLM = "llm", "LLM Usage"
+    EMBEDDING = "embedding", "Embedding Usage"
+
+
+class LLMUsageReason(models.TextChoices):
+    EXTRACTION = "extraction", "Extraction"
+    CHALLENGE = "challenge", "Challenge"
+    SUMMARIZE = "summarize", "Summarize"
+
+
 class Usage(BaseModel):
     id = models.UUIDField(
         primary_key=True,
@@ -26,7 +37,18 @@ class Usage(BaseModel):
     run_id = models.CharField(
         max_length=255, null=True, blank=True, db_comment="Identifier for the run"
     )
-    usage_type = models.CharField(max_length=255, db_comment="Type of usage")
+    usage_type = models.CharField(
+        max_length=255,
+        choices=UsageType.choices,
+        db_comment="Type of usage, either 'llm' or 'embedding'",
+    )
+    llm_usage_reason = models.CharField(
+        max_length=255,
+        choices=LLMUsageReason.choices,
+        null=True,
+        blank=True,
+        db_comment="Reason for LLM usage. Empty if usage_type is 'embedding'. ",
+    )
     model_name = models.CharField(max_length=255, db_comment="Name of the model used")
     embedding_tokens = models.IntegerField(
         db_comment="Number of tokens used for embedding"
@@ -38,9 +60,13 @@ class Usage(BaseModel):
         db_comment="Number of tokens used for the completion"
     )
     total_tokens = models.IntegerField(db_comment="Total number of tokens used")
+    cost_in_dollars = models.FloatField(db_comment="Total number of tokens used")
 
     def __str__(self):
         return str(self.id)
 
     class Meta:
         db_table = "token_usage"
+        indexes = [
+            models.Index(fields=["run_id"]),
+        ]