outline calls updated

ucbepic · Jan 11, 2025 · 636f3f3 · 636f3f3
1 parent 0b236b0
commit 636f3f3
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/docetl/operations/utils/api.py b/docetl/operations/utils/api.py
@@ -561,6 +561,19 @@ def _call_llm_with_cache(
         messages = truncate_messages(messages, model)
 
         self.runner.rate_limiter.try_acquire("llm_call", weight=1)
+        if self._is_outlines_model(model):
+            model_path = self._get_model_path(model)
+            return self.outlines_backend.process_messages(
+                model_path=model_path,
+                messages=[
+                    {
+                        "role": "system",
+                        "content": system_prompt,
+                    },
+                ] + messages,
+                output_schema=output_schema
+            )
+        # Handle other models through LiteLLM
         if tools is not None:
             try:
                 response = completion(