Add json model for custom llm endpoints

Author: nik

adala/runtimes/_litellm.py

@@ -174,6 +174,7 @@ def init_runtime(self) -> "Runtime":
             raise ValueError(
                 f'Failed to check availability of requested model "{self.model}": {e}'
             )
+
         return self
 
     def get_llm_response(self, messages: List[Dict[str, str]]) -> str:
@@ -343,6 +344,7 @@ def init_runtime(self) -> "Runtime":
             raise ValueError(
                 f'Failed to check availability of requested model "{self.model}": {e}'
             )
+
         return self
 
     @field_validator("concurrency", mode="before")
@@ -355,6 +357,10 @@ def check_concurrency(cls, value) -> int:
             )
         return value
 
+    @property
+    def is_custom_openai_endpoint(self) -> bool:
+        return self.model.startswith("openai/") and self.model_extra.get("base_url")
+
     async def batch_to_batch(
         self,
         batch: InternalDataFrame,
@@ -383,6 +389,12 @@ async def batch_to_batch(
         ).tolist()
 
         retries = AsyncRetrying(**RETRY_POLICY)
+        if self.is_custom_openai_endpoint:
+            # TODO: most of the custom openai endpoints do not support tools mode but json mode
+            # we should make it more performant by not creating instructor client on every request
+            async_instructor_client = instructor.from_litellm(
+                litellm.acompletion, mode=instructor.Mode.JSON
+            )
 
         tasks = [
             asyncio.ensure_future(

adala/skills/collection/entity_extraction.py

@@ -267,7 +267,7 @@ def extract_indices(self, df):
         input_field_name = self._get_input_field_name()
         output_field_name = self._get_output_field_name()
         for i, row in df.iterrows():
-            if row.get('_adala_error'):
+            if row.get("_adala_error"):
                 logger.warning(f"Error in row {i}: {row['_adala_message']}")
                 continue
             text = row[input_field_name]