Added saving of the trace data

src/architectures.py

@@ -12,8 +12,9 @@ import traceback
 
 from abc import ABC, abstractmethod
 from enum import Enum
+from huggingface_hub import Repository
 from time import time
-from typing import List, Optional
+from typing import List, Optional, Dict
 from better_profanity import profanity
 
 from src.common import config_dir, data_dir, hf_api_token, escape_dollars
@@ -62,6 +63,9 @@ class ArchitectureRequest:
             md += f"\n  - {r}"
         return escape_dollars(md)
 
+    def as_dict(self) -> Dict:
+        return {'request_evolution': self._request, 'response_evolution': self._response}
+
 
 class ArchitectureTraceOutcome(Enum):
     """
@@ -115,6 +119,16 @@ class ArchitectureTraceStep:
         md += f"  - **Outcome**: {outcome}"
         return escape_dollars(md)
 
+    def as_dict(self) -> Dict:
+        return {
+            'name': self.name,
+            'start_ms': self.start_ms,
+            'end_ms': self.end_ms,
+            'outcome': str(self.outcome),
+            'exception': "" if self._exception is None else f"{self._exception}",
+            'early_exit_message': "" if self.early_exit_message is None else self.early_exit_message
+        }
+
 
 class ArchitectureTrace:
     """
@@ -143,6 +157,9 @@ class ArchitectureTrace:
         md = '  \n'.join([s.as_markdown() for s in self.steps])
         return md
 
+    def as_dict(self) -> Dict:
+        return {'steps': [s.as_dict() for s in self.steps]}
+
 
 class ArchitectureComponent(ABC):
     description = "Components should override a description"
@@ -177,6 +194,8 @@ class Architecture:
     crash the system.
     """
     architectures = None
+    save_repo = None
+    save_repo_load_error = False
 
     @classmethod
     def load_architectures(cls, force_reload: bool = False) -> None:
@@ -221,6 +240,46 @@ class Architecture:
                 return a
         raise ValueError(f"Could not find an architecture named {name}")
 
+    def append_and_save_data_as_json(self, data: Dict):
+        repo_url = "https://huggingface.co/datasets/alfraser/llm-arch-trace"
+        trace_dir = "trace"
+        trace_file_name = "trace.json"
+        data_file = os.path.join(trace_dir, trace_file_name)
+        if Architecture.save_repo is None and not Architecture.save_repo_load_error:
+            try:
+                hf_write_token = hf_api_token(write=True)
+                Architecture.save_repo = Repository(local_dir=trace_dir, clone_from=repo_url, token=hf_write_token)
+            except Exception as err:
+                Architecture.save_repo_load_error = True
+                print(f"Error connecting to the save repo {err} - persistence now disabled")
+
+        if Architecture.save_repo is not None:
+            with open(data_file, 'r') as f:
+                test_json = json.load(f)
+            test_json['tests'].append(data)
+            with open(data_file, 'w') as f:
+                json.dump(test_json, f, indent=2)
+            Architecture.save_repo.push_to_hub()
+
+    def attempt_request_and_trace_save(self, request: ArchitectureRequest, trace: ArchitectureTrace,
+                                       trace_tags: List[str] = None) -> None:
+        """
+        Attempt to save a request and trace pair to a json store on huggingface datasets
+        Catch any errors and simply print as non-fatal to functional flow
+        """
+        try:
+            if trace_tags is None:
+                trace_tags = []
+            save_dict = {
+                'architecture': self.name,
+                'request': request.as_dict(),
+                'trace': trace.as_dict(),
+                'trace_tags': trace_tags
+            }
+            self.append_and_save_data_as_json(save_dict)
+        except Exception as err:
+            print(f"Request / trace save failed {err}")
+
     def __init__(self,
                  name: str,
                  description: str,
@@ -235,7 +294,7 @@ class Architecture:
         self.exception_text = exception_text
         self.no_response_text = no_response_text
 
-    def __call__(self, request: ArchitectureRequest) -> ArchitectureTrace:
+    def __call__(self, request: ArchitectureRequest, trace_tags: List[str] = None) -> ArchitectureTrace:
         """
         The main entry point to call the pipeline. Passes the request through each pipeline step
         in sequence, allowing them to amend the request or early exit the processing. Also captures
@@ -260,8 +319,7 @@ class Architecture:
                 trace.steps[-1].exception = err
                 traceback.print_exc()
                 break
-        # TODO - save the request / response
-        # TODO - save the trace
+        self.attempt_request_and_trace_save(request, trace, trace_tags)
         return trace
 
 
@@ -414,3 +472,10 @@ class ResponseTrimmer(ArchitectureComponent):
         return f"Regexes: {self.regex_display}"
 
 
+if __name__ == "__main__":
+    req = ArchitectureRequest("Testing")
+    a = Architecture.get_architecture("1. Baseline LLM")
+    a(req)
+    print("Hold")
+
+