From 19ba1a3099772fe608b1f6513ae179299cb1e885 Mon Sep 17 00:00:00 2001
From: Simon Willison <swillison@gmail.com>
Date: Tue, 11 Jul 2023 15:14:52 -0700
Subject: [PATCH] Improvements to Response.fake() plus __repr__ for some
 classes

---
 llm/models.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/llm/models.py b/llm/models.py
index 4a4f6b5..5f28195 100644
--- a/llm/models.py
+++ b/llm/models.py
@@ -150,7 +150,7 @@ class Response(ABC):
         db["responses"].insert(response)
 
     @classmethod
-    def fake(cls, model, prompt, system, response):
+    def fake(cls, model: "Model", prompt: str, system: str, response: str):
         "Utility method to help with writing tests"
         response_obj = cls(
             model=model,
@@ -182,12 +182,17 @@ class Response(ABC):
             stream=False,
         )
         response.id = row["id"]
-        response._prompt_json = json.loads(row["prompt_json"])
-        response.response_json = json.loads(row["response_json"])
+        response._prompt_json = json.loads(row["prompt_json"] or "null")
+        response.response_json = json.loads(row["response_json"] or "null")
         response._done = True
         response._chunks = [row["response"]]
         return response
 
+    def __repr__(self):
+        return "<Response prompt='{}' text='{}'>".format(
+            self.prompt.prompt, self.text()
+        )
+
 
 class Options(BaseModel):
     model_config = ConfigDict(extra="forbid")
@@ -258,6 +263,9 @@ class Model(ABC):
     def __str__(self) -> str:
         return "{}: {}".format(self.__class__.__name__, self.model_id)
 
+    def __repr__(self):
+        return "<Model '{}'>".format(self.model_id)
+
 
 @dataclass
 class ModelWithAliases: