kiln-ai 0.19.0__py3-none-any.whl → 0.21.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of kiln-ai might be problematic. Click here for more details.

Files changed (158) hide show
  1. kiln_ai/adapters/__init__.py +8 -2
  2. kiln_ai/adapters/adapter_registry.py +43 -208
  3. kiln_ai/adapters/chat/chat_formatter.py +8 -12
  4. kiln_ai/adapters/chat/test_chat_formatter.py +6 -2
  5. kiln_ai/adapters/chunkers/__init__.py +13 -0
  6. kiln_ai/adapters/chunkers/base_chunker.py +42 -0
  7. kiln_ai/adapters/chunkers/chunker_registry.py +16 -0
  8. kiln_ai/adapters/chunkers/fixed_window_chunker.py +39 -0
  9. kiln_ai/adapters/chunkers/helpers.py +23 -0
  10. kiln_ai/adapters/chunkers/test_base_chunker.py +63 -0
  11. kiln_ai/adapters/chunkers/test_chunker_registry.py +28 -0
  12. kiln_ai/adapters/chunkers/test_fixed_window_chunker.py +346 -0
  13. kiln_ai/adapters/chunkers/test_helpers.py +75 -0
  14. kiln_ai/adapters/data_gen/test_data_gen_task.py +9 -3
  15. kiln_ai/adapters/docker_model_runner_tools.py +119 -0
  16. kiln_ai/adapters/embedding/__init__.py +0 -0
  17. kiln_ai/adapters/embedding/base_embedding_adapter.py +44 -0
  18. kiln_ai/adapters/embedding/embedding_registry.py +32 -0
  19. kiln_ai/adapters/embedding/litellm_embedding_adapter.py +199 -0
  20. kiln_ai/adapters/embedding/test_base_embedding_adapter.py +283 -0
  21. kiln_ai/adapters/embedding/test_embedding_registry.py +166 -0
  22. kiln_ai/adapters/embedding/test_litellm_embedding_adapter.py +1149 -0
  23. kiln_ai/adapters/eval/base_eval.py +2 -2
  24. kiln_ai/adapters/eval/eval_runner.py +9 -3
  25. kiln_ai/adapters/eval/g_eval.py +2 -2
  26. kiln_ai/adapters/eval/test_base_eval.py +2 -4
  27. kiln_ai/adapters/eval/test_g_eval.py +4 -5
  28. kiln_ai/adapters/extractors/__init__.py +18 -0
  29. kiln_ai/adapters/extractors/base_extractor.py +72 -0
  30. kiln_ai/adapters/extractors/encoding.py +20 -0
  31. kiln_ai/adapters/extractors/extractor_registry.py +44 -0
  32. kiln_ai/adapters/extractors/extractor_runner.py +112 -0
  33. kiln_ai/adapters/extractors/litellm_extractor.py +386 -0
  34. kiln_ai/adapters/extractors/test_base_extractor.py +244 -0
  35. kiln_ai/adapters/extractors/test_encoding.py +54 -0
  36. kiln_ai/adapters/extractors/test_extractor_registry.py +181 -0
  37. kiln_ai/adapters/extractors/test_extractor_runner.py +181 -0
  38. kiln_ai/adapters/extractors/test_litellm_extractor.py +1192 -0
  39. kiln_ai/adapters/fine_tune/__init__.py +1 -1
  40. kiln_ai/adapters/fine_tune/openai_finetune.py +14 -4
  41. kiln_ai/adapters/fine_tune/test_dataset_formatter.py +2 -2
  42. kiln_ai/adapters/fine_tune/test_fireworks_tinetune.py +2 -6
  43. kiln_ai/adapters/fine_tune/test_openai_finetune.py +108 -111
  44. kiln_ai/adapters/fine_tune/test_together_finetune.py +2 -6
  45. kiln_ai/adapters/ml_embedding_model_list.py +192 -0
  46. kiln_ai/adapters/ml_model_list.py +761 -37
  47. kiln_ai/adapters/model_adapters/base_adapter.py +51 -21
  48. kiln_ai/adapters/model_adapters/litellm_adapter.py +380 -138
  49. kiln_ai/adapters/model_adapters/test_base_adapter.py +193 -17
  50. kiln_ai/adapters/model_adapters/test_litellm_adapter.py +407 -2
  51. kiln_ai/adapters/model_adapters/test_litellm_adapter_tools.py +1103 -0
  52. kiln_ai/adapters/model_adapters/test_saving_adapter_results.py +5 -5
  53. kiln_ai/adapters/model_adapters/test_structured_output.py +113 -5
  54. kiln_ai/adapters/ollama_tools.py +69 -12
  55. kiln_ai/adapters/parsers/__init__.py +1 -1
  56. kiln_ai/adapters/provider_tools.py +205 -47
  57. kiln_ai/adapters/rag/deduplication.py +49 -0
  58. kiln_ai/adapters/rag/progress.py +252 -0
  59. kiln_ai/adapters/rag/rag_runners.py +844 -0
  60. kiln_ai/adapters/rag/test_deduplication.py +195 -0
  61. kiln_ai/adapters/rag/test_progress.py +785 -0
  62. kiln_ai/adapters/rag/test_rag_runners.py +2376 -0
  63. kiln_ai/adapters/remote_config.py +80 -8
  64. kiln_ai/adapters/repair/test_repair_task.py +12 -9
  65. kiln_ai/adapters/run_output.py +3 -0
  66. kiln_ai/adapters/test_adapter_registry.py +657 -85
  67. kiln_ai/adapters/test_docker_model_runner_tools.py +305 -0
  68. kiln_ai/adapters/test_ml_embedding_model_list.py +429 -0
  69. kiln_ai/adapters/test_ml_model_list.py +251 -1
  70. kiln_ai/adapters/test_ollama_tools.py +340 -1
  71. kiln_ai/adapters/test_prompt_adaptors.py +13 -6
  72. kiln_ai/adapters/test_prompt_builders.py +1 -1
  73. kiln_ai/adapters/test_provider_tools.py +254 -8
  74. kiln_ai/adapters/test_remote_config.py +651 -58
  75. kiln_ai/adapters/vector_store/__init__.py +1 -0
  76. kiln_ai/adapters/vector_store/base_vector_store_adapter.py +83 -0
  77. kiln_ai/adapters/vector_store/lancedb_adapter.py +389 -0
  78. kiln_ai/adapters/vector_store/test_base_vector_store.py +160 -0
  79. kiln_ai/adapters/vector_store/test_lancedb_adapter.py +1841 -0
  80. kiln_ai/adapters/vector_store/test_vector_store_registry.py +199 -0
  81. kiln_ai/adapters/vector_store/vector_store_registry.py +33 -0
  82. kiln_ai/datamodel/__init__.py +39 -34
  83. kiln_ai/datamodel/basemodel.py +170 -1
  84. kiln_ai/datamodel/chunk.py +158 -0
  85. kiln_ai/datamodel/datamodel_enums.py +28 -0
  86. kiln_ai/datamodel/embedding.py +64 -0
  87. kiln_ai/datamodel/eval.py +1 -1
  88. kiln_ai/datamodel/external_tool_server.py +298 -0
  89. kiln_ai/datamodel/extraction.py +303 -0
  90. kiln_ai/datamodel/json_schema.py +25 -10
  91. kiln_ai/datamodel/project.py +40 -1
  92. kiln_ai/datamodel/rag.py +79 -0
  93. kiln_ai/datamodel/registry.py +0 -15
  94. kiln_ai/datamodel/run_config.py +62 -0
  95. kiln_ai/datamodel/task.py +2 -77
  96. kiln_ai/datamodel/task_output.py +6 -1
  97. kiln_ai/datamodel/task_run.py +41 -0
  98. kiln_ai/datamodel/test_attachment.py +649 -0
  99. kiln_ai/datamodel/test_basemodel.py +4 -4
  100. kiln_ai/datamodel/test_chunk_models.py +317 -0
  101. kiln_ai/datamodel/test_dataset_split.py +1 -1
  102. kiln_ai/datamodel/test_embedding_models.py +448 -0
  103. kiln_ai/datamodel/test_eval_model.py +6 -6
  104. kiln_ai/datamodel/test_example_models.py +175 -0
  105. kiln_ai/datamodel/test_external_tool_server.py +691 -0
  106. kiln_ai/datamodel/test_extraction_chunk.py +206 -0
  107. kiln_ai/datamodel/test_extraction_model.py +470 -0
  108. kiln_ai/datamodel/test_rag.py +641 -0
  109. kiln_ai/datamodel/test_registry.py +8 -3
  110. kiln_ai/datamodel/test_task.py +15 -47
  111. kiln_ai/datamodel/test_tool_id.py +320 -0
  112. kiln_ai/datamodel/test_vector_store.py +320 -0
  113. kiln_ai/datamodel/tool_id.py +105 -0
  114. kiln_ai/datamodel/vector_store.py +141 -0
  115. kiln_ai/tools/__init__.py +8 -0
  116. kiln_ai/tools/base_tool.py +82 -0
  117. kiln_ai/tools/built_in_tools/__init__.py +13 -0
  118. kiln_ai/tools/built_in_tools/math_tools.py +124 -0
  119. kiln_ai/tools/built_in_tools/test_math_tools.py +204 -0
  120. kiln_ai/tools/mcp_server_tool.py +95 -0
  121. kiln_ai/tools/mcp_session_manager.py +246 -0
  122. kiln_ai/tools/rag_tools.py +157 -0
  123. kiln_ai/tools/test_base_tools.py +199 -0
  124. kiln_ai/tools/test_mcp_server_tool.py +457 -0
  125. kiln_ai/tools/test_mcp_session_manager.py +1585 -0
  126. kiln_ai/tools/test_rag_tools.py +848 -0
  127. kiln_ai/tools/test_tool_registry.py +562 -0
  128. kiln_ai/tools/tool_registry.py +85 -0
  129. kiln_ai/utils/__init__.py +3 -0
  130. kiln_ai/utils/async_job_runner.py +62 -17
  131. kiln_ai/utils/config.py +24 -2
  132. kiln_ai/utils/env.py +15 -0
  133. kiln_ai/utils/filesystem.py +14 -0
  134. kiln_ai/utils/filesystem_cache.py +60 -0
  135. kiln_ai/utils/litellm.py +94 -0
  136. kiln_ai/utils/lock.py +100 -0
  137. kiln_ai/utils/mime_type.py +38 -0
  138. kiln_ai/utils/open_ai_types.py +94 -0
  139. kiln_ai/utils/pdf_utils.py +38 -0
  140. kiln_ai/utils/project_utils.py +17 -0
  141. kiln_ai/utils/test_async_job_runner.py +151 -35
  142. kiln_ai/utils/test_config.py +138 -1
  143. kiln_ai/utils/test_env.py +142 -0
  144. kiln_ai/utils/test_filesystem_cache.py +316 -0
  145. kiln_ai/utils/test_litellm.py +206 -0
  146. kiln_ai/utils/test_lock.py +185 -0
  147. kiln_ai/utils/test_mime_type.py +66 -0
  148. kiln_ai/utils/test_open_ai_types.py +131 -0
  149. kiln_ai/utils/test_pdf_utils.py +73 -0
  150. kiln_ai/utils/test_uuid.py +111 -0
  151. kiln_ai/utils/test_validation.py +524 -0
  152. kiln_ai/utils/uuid.py +9 -0
  153. kiln_ai/utils/validation.py +90 -0
  154. {kiln_ai-0.19.0.dist-info → kiln_ai-0.21.0.dist-info}/METADATA +12 -5
  155. kiln_ai-0.21.0.dist-info/RECORD +211 -0
  156. kiln_ai-0.19.0.dist-info/RECORD +0 -115
  157. {kiln_ai-0.19.0.dist-info → kiln_ai-0.21.0.dist-info}/WHEEL +0 -0
  158. {kiln_ai-0.19.0.dist-info → kiln_ai-0.21.0.dist-info}/licenses/LICENSE.txt +0 -0
@@ -25,8 +25,11 @@ from kiln_ai.datamodel import (
25
25
  )
26
26
  from kiln_ai.datamodel.datamodel_enums import ChatStrategy
27
27
  from kiln_ai.datamodel.json_schema import validate_schema_with_value_error
28
- from kiln_ai.datamodel.task import RunConfig
28
+ from kiln_ai.datamodel.task import RunConfigProperties
29
+ from kiln_ai.tools import KilnToolInterface
30
+ from kiln_ai.tools.tool_registry import tool_from_id
29
31
  from kiln_ai.utils.config import Config
32
+ from kiln_ai.utils.open_ai_types import ChatCompletionMessageParam
30
33
 
31
34
 
32
35
  @dataclass
@@ -48,33 +51,24 @@ class BaseAdapter(metaclass=ABCMeta):
48
51
  This abstract class provides the foundation for implementing model-specific adapters
49
52
  that can process tasks with structured or unstructured inputs/outputs. It handles
50
53
  input/output validation, prompt building, and run tracking.
51
-
52
- Attributes:
53
- prompt_builder (BasePromptBuilder): Builder for constructing prompts for the model
54
- kiln_task (Task): The task configuration and metadata
55
- output_schema (dict | None): JSON schema for validating structured outputs
56
- input_schema (dict | None): JSON schema for validating structured inputs
57
54
  """
58
55
 
59
56
  def __init__(
60
57
  self,
61
- run_config: RunConfig,
58
+ task: Task,
59
+ run_config: RunConfigProperties,
62
60
  config: AdapterConfig | None = None,
63
61
  ):
62
+ self.task = task
64
63
  self.run_config = run_config
65
64
  self.update_run_config_unknown_structured_output_mode()
66
- self.prompt_builder = prompt_builder_from_id(
67
- run_config.prompt_id, run_config.task
68
- )
65
+ self.prompt_builder = prompt_builder_from_id(run_config.prompt_id, task)
69
66
  self._model_provider: KilnModelProvider | None = None
70
67
 
71
- self.output_schema = self.task().output_json_schema
72
- self.input_schema = self.task().input_json_schema
68
+ self.output_schema = task.output_json_schema
69
+ self.input_schema = task.input_json_schema
73
70
  self.base_adapter_config = config or AdapterConfig()
74
71
 
75
- def task(self) -> Task:
76
- return self.run_config.task
77
-
78
72
  def model_provider(self) -> KilnModelProvider:
79
73
  """
80
74
  Lazy load the model provider for this adapter.
@@ -152,7 +146,11 @@ class BaseAdapter(metaclass=ABCMeta):
152
146
  f"response is not a string for non-structured task: {parsed_output.output}"
153
147
  )
154
148
 
155
- # Validate reasoning content is present (if reasoning)
149
+ # Validate reasoning content is present and required
150
+ # We don't require reasoning when using tools as models tend not to return any on the final turn (both Sonnet and Gemini).
151
+ trace_has_toolcalls = parsed_output.trace is not None and any(
152
+ message.get("role", None) == "tool" for message in parsed_output.trace
153
+ )
156
154
  if (
157
155
  provider.reasoning_capable
158
156
  and (
@@ -163,19 +161,22 @@ class BaseAdapter(metaclass=ABCMeta):
163
161
  provider.reasoning_optional_for_structured_output
164
162
  and self.has_structured_output()
165
163
  )
164
+ and not (trace_has_toolcalls)
166
165
  ):
167
166
  raise RuntimeError(
168
167
  "Reasoning is required for this model, but no reasoning was returned."
169
168
  )
170
169
 
171
170
  # Generate the run and output
172
- run = self.generate_run(input, input_source, parsed_output, usage)
171
+ run = self.generate_run(
172
+ input, input_source, parsed_output, usage, run_output.trace
173
+ )
173
174
 
174
175
  # Save the run if configured to do so, and we have a path to save to
175
176
  if (
176
177
  self.base_adapter_config.allow_saving
177
178
  and Config.shared().autosave_runs
178
- and self.task().path is not None
179
+ and self.task.path is not None
179
180
  ):
180
181
  run.save_to_file()
181
182
  else:
@@ -261,6 +262,7 @@ class BaseAdapter(metaclass=ABCMeta):
261
262
  input_source: DataSource | None,
262
263
  run_output: RunOutput,
263
264
  usage: Usage | None = None,
265
+ trace: list[ChatCompletionMessageParam] | None = None,
264
266
  ) -> TaskRun:
265
267
  # Convert input and output to JSON strings if they are dictionaries
266
268
  input_str = (
@@ -280,7 +282,7 @@ class BaseAdapter(metaclass=ABCMeta):
280
282
  )
281
283
 
282
284
  new_task_run = TaskRun(
283
- parent=self.task(),
285
+ parent=self.task,
284
286
  input=input_str,
285
287
  input_source=input_source,
286
288
  output=TaskOutput(
@@ -289,11 +291,13 @@ class BaseAdapter(metaclass=ABCMeta):
289
291
  source=DataSource(
290
292
  type=DataSourceType.synthetic,
291
293
  properties=self._properties_for_task_output(),
294
+ run_config=self.run_config,
292
295
  ),
293
296
  ),
294
297
  intermediate_outputs=run_output.intermediate_outputs,
295
298
  tags=self.base_adapter_config.default_tags or [],
296
299
  usage=usage,
300
+ trace=trace,
297
301
  )
298
302
 
299
303
  return new_task_run
@@ -301,8 +305,10 @@ class BaseAdapter(metaclass=ABCMeta):
301
305
  def _properties_for_task_output(self) -> Dict[str, str | int | float]:
302
306
  props = {}
303
307
 
304
- # adapter info
305
308
  props["adapter_name"] = self.adapter_name()
309
+
310
+ # Legacy properties where we save the run_config details into custom properties.
311
+ # These are now also be saved in the run_config field.
306
312
  props["model_name"] = self.run_config.model_name
307
313
  props["model_provider"] = self.run_config.model_provider_name
308
314
  props["prompt_id"] = self.run_config.prompt_id
@@ -325,3 +331,27 @@ class BaseAdapter(metaclass=ABCMeta):
325
331
  )
326
332
  new_run_config.structured_output_mode = structured_output_mode
327
333
  self.run_config = new_run_config
334
+
335
+ async def available_tools(self) -> list[KilnToolInterface]:
336
+ tool_config = self.run_config.tools_config
337
+ if tool_config is None or tool_config.tools is None:
338
+ return []
339
+
340
+ project = self.task.parent_project()
341
+ if project is None:
342
+ raise ValueError("Task must have a parent project to resolve tools")
343
+
344
+ project_id = project.id
345
+ if project_id is None:
346
+ raise ValueError("Project must have an ID to resolve tools")
347
+
348
+ tools = [tool_from_id(tool_id, self.task) for tool_id in tool_config.tools]
349
+
350
+ # Check each tool has a unique name
351
+ tool_names = [await tool.name() for tool in tools]
352
+ if len(tool_names) != len(set(tool_names)):
353
+ raise ValueError(
354
+ "Each tool must have a unique name. Either de-select the duplicate tools, or modify their names to describe their unique purpose. Model will struggle if tools do not have descriptive names and tool execution will be undefined."
355
+ )
356
+
357
+ return tools