add more info to the prompt history

aelaguiz · Mar 20, 2024 · 56d7b24 · 56d7b24
1 parent e94463c
commit 56d7b24
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/langdspy/prompt_runners.py b/langdspy/prompt_runners.py
@@ -39,13 +39,15 @@ def __init__(self, **kwargs):
 class PromptHistory(BaseModel):
     history: List[Dict[str, Any]] = Field(default_factory=list)
 
-    def add_entry(self, prompt, llm_response, parsed_output, success, timestamp):
+    def add_entry(self, llm, prompt, llm_response, parsed_output, success, start_time, end_time):
         self.history.append({
-            "prompt": prompt,
+            "duration_ms": round((end_time - start_time) * 1000),
+            "llm": llm,
             "llm_response": llm_response,
             "parsed_output": parsed_output,
+            "prompt": prompt,
             "success": success,
-            "timestamp": timestamp
+            "timestamp": end_time,
         })
 
 
@@ -101,6 +103,7 @@ def _invoke_with_retries(self, chain, input, max_tries=1, config: Optional[Runna
         formatted_prompt = None
 
         while max_tries >= 1:
+            start_time = time.time()
             try:
                 kwargs = {**self.model_kwargs, **self.kwargs}
                 # logger.debug(f"PromptRunner invoke with input {input} and kwargs {kwargs} and config {config}")
@@ -162,7 +165,8 @@ def _invoke_with_retries(self, chain, input, max_tries=1, config: Optional[Runna
                 logger.error(f"Output keys do not match expected output keys for prompt runner {self.template.__class__.__name__}")
                 validation = False
 
-            self.prompt_history.add_entry(formatted_prompt, res, parsed_output, validation, time.time())
+            end_time = time.time()
+            self.prompt_history.add_entry(config["llm"].model_name, formatted_prompt, res, parsed_output, validation, start_time, end_time)
 
             if validation:
                 # Transform and validate the outputs
@@ -263,4 +267,4 @@ def run_task():
 
 
         # logger.debug(f"MultiPromptRunner predictions: {self.predictions}")
-        return predictions
+        return predictions