Added dumb trim

Giskard-AI · Nov 19, 2024 · b09d266 · b09d266
1 parent 39c4fa9
commit b09d266
Showing 1 changed file with 17 additions and 0 deletions.
diff --git a/giskard/llm/client/litellm.py b/giskard/llm/client/litellm.py
@@ -22,6 +22,20 @@ def _get_response_format(format):
     return None
 
 
+def _json_trim(response_message: str):
+    # Dumb trim for when model response message in addition to the JSON response
+    if "{" not in response_message or "}" not in response_message:
+        raise ValueError("The model output doesn't contain any JSON")
+
+    json_start = response_message.index("{")
+    json_end = len(response_message) - response_message[::-1].index("}")
+
+    if json_start > json_end:
+        raise ValueError("The model output doesn't contain any JSON")
+
+    return response_message
+
+
 class LiteLLMClient(LLMClient):
     def __init__(self, model: str = "gpt-4o", completion_params: Optional[Dict[str, Any]] = None):
         """Initialize a LiteLLM completion client
@@ -77,4 +91,7 @@ def complete(
 
         response_message = completion.choices[0].message
 
+        if format is not None:
+            response_message = _json_trim(response_message)
+
         return ChatMessage(role=response_message.role, content=response_message.content)