add token cost.

data_processing.py

@@ -22,6 +22,7 @@ import google.generativeai as genai
 import urllib
 import pysolr
 import config # Import the config module to access remote host details
+import tiktoken
 
 from llm_prompts import (
     get_analysis_plan_prompt,
@@ -51,11 +52,11 @@ def llm_generate_analysis_plan_with_history(llm_model, natural_language_query, c
 
         if intent != 'search_list':
             print(f"API returned intent '{intent}' which is not 'search_list'. Aborting analysis.")
-            return None, None, None, intent
+            return None, None, None, intent, None, None, None
 
     except Exception as e:
         print(f"Warning: Could not retrieve dynamic search fields. Proceeding without them. Error: {e}")
-        return None, [], None, 'api_error'
+        return None, [], None, 'api_error', None, None, None
 
     core_name = search_name if search_name else 'news'
 
@@ -75,17 +76,21 @@ def llm_generate_analysis_plan_with_history(llm_model, natural_language_query, c
 
     try:
         response = llm_model.generate_content(prompt)
+        encoding = tiktoken.encoding_for_model("gpt-4")
+        input_token_count = len(encoding.encode(prompt))
+        output_token_count = len(encoding.encode(response.text))
+        total_token_count = (input_token_count if input_token_count is not None else 0) + (output_token_count if output_token_count is not None else 0)
         cleaned_text = re.sub(r'```json\s*|\s*```', '', response.text, flags=re.MULTILINE | re.DOTALL).strip()
         plan = json.loads(cleaned_text)
-        return plan, mapped_search_fields, core_name, intent
+        return plan, mapped_search_fields, core_name, intent, input_token_count, output_token_count, total_token_count
     except json.JSONDecodeError as e:
         raw_response_text = response.text if 'response' in locals() else 'N/A'
         print(f"Error decoding JSON from LLM response: {e}\nRaw Response:\n{raw_response_text}")
-        return None, mapped_search_fields, core_name, intent
+        return None, mapped_search_fields, core_name, intent, None, None, None
     except Exception as e:
         raw_response_text = response.text if 'response' in locals() else 'N/A'
         print(f"Error in llm_generate_analysis_plan_with_history: {e}\nRaw Response:\n{raw_response_text}")
-        return None, mapped_search_fields, core_name, intent
+        return None, mapped_search_fields, core_name, intent, None, None, None
 
 def execute_quantitative_query(solr_client, plan):
     """Executes the facet query to get aggregate data."""
@@ -150,12 +155,24 @@ def llm_synthesize_enriched_report_stream(llm_model, query, quantitative_data, q
     """
     prompt = get_synthesis_report_prompt(query, quantitative_data, qualitative_data, plan)
     try:
-        response_stream = llm_model.generate_content(prompt, stream=True)
+        response_stream = llm_model.generate_content(prompt, stream=True)        
+        response_text = ""
         for chunk in response_stream:
-            yield chunk.text
+            yield {"text": chunk.text, "tokens": None}
+            response_text += chunk.text
+        encoding = tiktoken.encoding_for_model("gpt-4")
+        input_token_count = len(encoding.encode(prompt))
+        output_token_count = len(encoding.encode(response_text))
+        total_token_count = (input_token_count if input_token_count is not None else 0) + (output_token_count if output_token_count is not None else 0)
+        tokens = {
+            "input": input_token_count,
+            "output": output_token_count,
+            "total": total_token_count,
+        }
+        yield {"text": None, "tokens": tokens}
     except Exception as e:
         print(f"Error in llm_synthesize_enriched_report_stream: {e}")
-        yield "Sorry, an error occurred while generating the report. Please check the logs for details."
+        yield {"text": "Sorry, an error occurred while generating the report. Please check the logs for details.", "tokens": None}
 
 def llm_generate_visualization_code(llm_model, query_context, facet_data):
     """Generates Python code for visualization based on query and data."""
@@ -163,12 +180,16 @@ def llm_generate_visualization_code(llm_model, query_context, facet_data):
     try:
         generation_config = genai.types.GenerationConfig(temperature=0)
         response = llm_model.generate_content(prompt, generation_config=generation_config)
+        encoding = tiktoken.encoding_for_model("gpt-4")
+        input_token_count = len(encoding.encode(prompt))
+        output_token_count = len(encoding.encode(response.text))
+        total_token_count = (input_token_count if input_token_count is not None else 0) + (output_token_count if output_token_count is not None else 0)
         code = re.sub(r'^```python\s*|```$', '', response.text, flags=re.MULTILINE)
-        return code
+        return code, input_token_count, output_token_count, total_token_count
     except Exception as e:
         raw_response_text = response.text if 'response' in locals() else 'N/A'
         print(f"Error in llm_generate_visualization_code: {e}\nRaw response: {raw_response_text}")
-        return None
+        return
 
 def execute_viz_code_and_get_path(viz_code, facet_data):
     """Executes visualization code and returns the path to the saved plot image."""