[#189] Add token information to WorkflowResult

labs/config/redis_client.py

@@ -10,11 +10,13 @@ class RedisVariable(Enum):
     BRANCH_NAME = "branch_name"
     CONTEXT = "context"
     EMBEDDINGS = "embeddings"
+    EMBEDDINGS_TOKENS = "embeddings_tokens"
     FILES_MODIFIED = "files_modified"
     ISSUE_BODY = "issue_body"
     ISSUE_NUMBER = "issue_number"
     ISSUE_TITLE = "issue_title"
     LLM_RESPONSE = "llm_response"
+    LLM_TOKENS = "llm_tokens"
     ORIGINAL_BRANCH_NAME = "original_branch_name"
     PRE_COMMIT_ERROR = "pre_commit_error"
     PROJECT = "project"

labs/core/admin.py

@@ -1,12 +1,10 @@
 from django.contrib import admin
 from django.urls import reverse
 
-from .forms import ProjectForm, EmbeddingModelFormSet, LLMModelFormSet
+from .forms import EmbeddingModelFormSet, LLMModelFormSet, ProjectForm
 from .mixins import JSONFormatterMixin
-from .models import (
-    EmbeddingModel, LLMModel, Project, Prompt, 
-    Variable, VectorizerModel, WorkflowResult
-)
+from .models import EmbeddingModel, LLMModel, Project, Prompt, Variable, VectorizerModel, WorkflowResult
+
 
 @admin.register(EmbeddingModel)
 class EmbeddingModelAdmin(admin.ModelAdmin):
@@ -165,8 +163,10 @@ class WorkflowResultAdmin(admin.ModelAdmin, JSONFormatterMixin):
         "embed_model",
         "prompt_model",
         "pretty_embeddings",
+        "embeddings_tokens",
         "pretty_context",
         "pretty_llm_response",
+        "llm_tokens",
         "pretty_modified_files",
         "pretty_pre_commit_error",
     )

labs/core/migrations/0010_workflowresult_embeddings_tokens.py

@@ -0,0 +1,17 @@
+# Generated by Django 5.1.5 on 2025-03-11 11:10
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("core", "0009_embeddingmodel_llmmodel_delete_model_and_more"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="workflowresult",
+            name="embeddings_tokens",
+            field=models.IntegerField(default=None, null=True),
+        ),
+    ]

labs/core/migrations/0011_workflowresult_prompt_tokens.py

@@ -0,0 +1,17 @@
+# Generated by Django 5.1.5 on 2025-03-11 11:42
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("core", "0010_workflowresult_embeddings_tokens"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="workflowresult",
+            name="prompt_tokens",
+            field=models.IntegerField(default=None, null=True),
+        ),
+    ]

labs/core/migrations/0012_remove_workflowresult_prompt_tokens_and_more.py

@@ -0,0 +1,21 @@
+# Generated by Django 5.1.7 on 2025-03-12 11:24
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("core", "0011_workflowresult_prompt_tokens"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="workflowresult",
+            name="prompt_tokens",
+        ),
+        migrations.AddField(
+            model_name="workflowresult",
+            name="llm_tokens",
+            field=models.IntegerField(default=None, null=True),
+        ),
+    ]

labs/core/models.py

@@ -26,10 +26,11 @@
 }
 
 vectorizer_model_class = {
-    "CHUNK_VECTORIZER": ChunkVectorizer, 
+    "CHUNK_VECTORIZER": ChunkVectorizer,
     "PYTHON_VECTORIZER": PythonVectorizer,
 }
 
+
 class ProviderEnum(Enum):
     NO_PROVIDER = "No provider"
     OPENAI = "OpenAI"
@@ -86,14 +87,14 @@ def _default_vectorizer_value_validation(self):
 
     def __str__(self) -> str:
         return self.name
-    
+
     class Meta:
         unique_together = ("provider", "name")
 
 
 class EmbeddingModel(models.Model):
     provider = models.CharField(choices=ProviderEnum.choices())
-    name = models.CharField(max_length=255,help_text="Ensure this Embedding exists and is downloaded.")
+    name = models.CharField(max_length=255, help_text="Ensure this Embedding exists and is downloaded.")
     active = models.BooleanField(default=True, help_text="Only one Embedding can be active.")
 
     @classmethod
@@ -133,10 +134,7 @@ class LLMModel(models.Model):
     name = models.CharField(max_length=255, help_text="Ensure this LLM exists and is downloaded.")
     active = models.BooleanField(default=True, help_text="Only one LLM can be active.")
     max_output_tokens = models.IntegerField(
-        null=True, 
-        blank=True, 
-        default=None,  
-        help_text="Leave blank for auto-detection, set only if required."
+        null=True, blank=True, default=None, help_text="Leave blank for auto-detection, set only if required."
     )
 
     @classmethod
@@ -208,7 +206,7 @@ def get_active_vectorizer(project_id: int) -> Vectorizer:
         vector_model = VectorizerModel.objects.filter(project__id=project_id).first()
         if not vector_model:
             raise ValueError("No vectorizer configured for this project.")
-        
+
         try:
             vec_class = vectorizer_model_class[vector_model.vectorizer_type]
         except KeyError:
@@ -230,8 +228,16 @@ class WorkflowResult(models.Model):
     embed_model = models.CharField(max_length=255, null=True)
     prompt_model = models.CharField(max_length=255, null=True)
     embeddings = models.TextField(null=True)
+    embeddings_tokens = models.IntegerField(
+        null=True,
+        default=None,
+    )
     context = models.TextField(null=True)
     llm_response = models.TextField(null=True)
+    llm_tokens = models.IntegerField(
+        null=True,
+        default=None,
+    )
     modified_files = models.TextField(null=True)
     pre_commit_error = models.TextField(null=True)
     created_at = models.DateTimeField(auto_now_add=True)

labs/embeddings/embedder.py

@@ -13,7 +13,8 @@
 @dataclass
 class Embeddings:
     model: str
-    embeddings: Union[List[Dict[str, Any]], List[List[int]]]
+    tokens: Optional[int]
+    embeddings: List[List[float]]
     model_config: Optional[Dict[str, Any]] = None
 
 
@@ -26,20 +27,16 @@ def embed(self, prompt, *args, **kwargs) -> Embeddings:
 
     def retrieve_files_path(
         self,
-        query: str,
+        embedded_prompt: Embeddings,
         project_id: int,
         similarity_threshold: float = settings.EMBEDDINGS_SIMILARITY_THRESHOLD,
         max_results: int = settings.EMBEDDINGS_MAX_RESULTS,
     ) -> List[str]:
-        query = query.replace("\n", "")
-        embedded_query = self.embed(prompt=query).embeddings
-        if not embedded_query:
-            raise ValueError(f"No embeddings found with the given {query=} with {similarity_threshold=}")
-
         files_path = (
             Embedding.objects.filter(project__id=project_id)
-            .values("file_path")  # the combination of values and annotate, is the Django way of making a group by
-            .annotate(distance=Min(CosineDistance("embedding", embedded_query[0])))
+            # the combination of values and annotate, is the Django way of making a group by
+            .values("file_path")
+            .annotate(distance=Min(CosineDistance("embedding", embedded_prompt.embeddings[0])))
             .order_by("distance")
             .values_list("file_path", flat=True)
         )[:max_results]

labs/embeddings/gemini.py

@@ -1,32 +1,37 @@
 import os
-import google.generativeai as genai
+
 from embeddings.embedder import Embeddings
+from google import genai
+
 
 class GeminiEmbedder:
     def __init__(self, model):
         self._model_name = model.name
         api_key = os.environ.get("GEMINI_API_KEY")
-        genai.configure(api_key=api_key)
+        self._client = genai.Client(api_key=api_key)
 
     def embed(self, prompt: str, *args, **kwargs) -> Embeddings:
         try:
-            result = genai.embed_content(
+            result = self._client.models.embed_content(
                 model=self._model_name,
-                content=prompt,
-                *args, 
+                contents=[prompt],
+                *args,
                 **kwargs,
             )
 
-            emb = result.get("embedding")
-            if isinstance(emb, list) and all(isinstance(e, list) for e in emb):
-                flat_vectors = emb
-            else:
-                flat_vectors = [emb]
-
+            assert result.embeddings
+
+            embeddings = []
+
+            for content_embedding in result.embeddings:
+                if content_embedding.values:
+                    embeddings.append(content_embedding.values)
+
             return Embeddings(
                 model=self._model_name,
-                model_config=result.get("model_config", {}),
-                embeddings=flat_vectors
+                model_config=result.model_config,
+                embeddings=embeddings,
+                tokens=None,
             )
 
         except Exception as e:

labs/embeddings/ollama.py

@@ -10,4 +10,4 @@ def __init__(self, model):
 
     def embed(self, prompt, *args, **kwargs) -> Embeddings:
         result = self._client.embed(self._model_name, prompt, *args, **kwargs)
-        return Embeddings(model=result["model"], embeddings=result["embeddings"])
+        return Embeddings(model=result["model"], embeddings=result["embeddings"], tokens=None)

labs/embeddings/openai.py

@@ -1,4 +1,5 @@
 import os
+
 import openai
 from embeddings.embedder import Embeddings
 from litellm import embedding
@@ -11,8 +12,13 @@ def __init__(self, model):
 
     def embed(self, prompt, *args, **kwargs) -> Embeddings:
         result = embedding(model=self._model_name, input=prompt, *args, **kwargs)
+
+        assert result.model is not None
+        assert result.usage is not None
+
         return Embeddings(
             model=result.model,
             model_config=result.model_config,
             embeddings=[data["embedding"] for data in result.data],
+            tokens=result.usage.total_tokens,
         )

labs/llm/anthropic.py

@@ -20,7 +20,7 @@ def completion_without_proxy(
         messages: List[Dict[str, str]],
         *args,
         **kwargs,
-    ) -> Tuple[str, Dict[str, Any]]:
+    ) -> Tuple[str, Dict[str, Any], int]:
         system_prompt = "\n".join([message["content"] for message in messages if message["role"] == "system"])
         user_messages = [message for message in messages if message["role"] == "user"]
 
@@ -32,9 +32,11 @@ def completion_without_proxy(
                 max_tokens=self._model_max_output_tokens,
             )
 
+            tokens = response.usage.input_tokens + response.usage.output_tokens
+
             response_steps = self.response_to_steps(response)
 
-            return self._model_name, {"choices": [{"message": {"content": response_steps}}]}
+            return self._model_name, {"choices": [{"message": {"content": response_steps}}]}, tokens
         except Exception as e:
             raise RuntimeError(f"Anthropic API call failed: {e}") from e
 

labs/llm/gemini.py

@@ -1,33 +1,38 @@
-import os
 import json
-from typing import List, Dict, Tuple, Any
+import logging
+import os
+from typing import Any, Dict, List, Tuple, cast
+
+from google import genai
+from google.genai import types
 
-import google.generativeai as genai
+logger = logging.getLogger(__name__)
 
 
 class GeminiRequester:
     def __init__(self, model):
         self._model_name = model.name
         api_key = os.environ.get("GEMINI_API_KEY")
-        genai.configure(api_key=api_key)
-        self.generative_model = genai.GenerativeModel(self._model_name)
-        self.generation_config = genai.GenerationConfig(response_mime_type="application/json")
+        self._client = genai.Client(api_key=api_key)
 
     def completion_without_proxy(
         self,
         messages: List[Dict[str, str]],
         *args,
         **kwargs,
-    ) -> Tuple[str, Dict[str, Any]]:
+    ) -> Tuple[str, Dict[str, Any], int]:
         try:
-            gemini_response = self.generative_model.generate_content(
+            gemini_response = self._client.models.generate_content(
+                model=self._model_name,
                 contents=json.dumps(messages),
-                generation_config=self.generation_config,
+                config=types.GenerateContentConfig(response_mime_type="application/json"),
                 *args,
                 **kwargs,
             )
-            return self._model_name, {
-                "choices": [{"message": {"content": gemini_response.text}}]
-            }
+
+            assert gemini_response.usage_metadata
+            tokens = cast(int, gemini_response.usage_metadata.total_token_count)
+
+            return self._model_name, {"choices": [{"message": {"content": gemini_response.text}}]}, tokens
         except Exception as e:
             raise RuntimeError(f"Gemini API call failed: {e}") from e

labs/llm/openai.py

@@ -1,20 +1,25 @@
 import os
-from typing import Tuple, Dict, Any
+from typing import Any, Dict, Tuple
 
 import openai
 from litellm import completion
+from litellm.types.utils import ModelResponse
 
 
 class OpenAIRequester:
     def __init__(self, model):
         self._model_name = model.name
         openai.api_key = os.environ.get("OPENAI_API_KEY")
 
-    def completion_without_proxy(self, messages, *args, **kwargs) -> Tuple[str, Dict[str, Any]]:
-        return self._model_name, completion(
+    def completion_without_proxy(self, messages, *args, **kwargs) -> Tuple[str, Dict[str, Any], int]:
+        response: ModelResponse = completion(
             model=self._model_name,
             messages=messages,
             response_format={"type": "json_object"},
             *args,
             **kwargs,
         )
+
+        tokens: int = response["usage"]["total_tokens"]
+
+        return self._model_name, response, tokens

labs/llm/requester.py

@@ -1,9 +1,11 @@
-from typing import Dict, List
+from typing import Any, Dict, List, Tuple
 
 
 class Requester:
     def __init__(self, requester, *args, **kwargs):
         self.requester = requester(*args, **kwargs)
 
-    def completion_without_proxy(self, messages: List[Dict[str, str]], *args, **kwargs):
+    def completion_without_proxy(
+        self, messages: List[Dict[str, str]], *args, **kwargs
+    ) -> Tuple[str, Dict[str, Any], int]:
         return self.requester.completion_without_proxy(messages, *args, **kwargs)