clarifai 11.4.10__tar.gz → 11.5.1__tar.gz

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.
Files changed (136) hide show
  1. {clarifai-11.4.10/clarifai.egg-info → clarifai-11.5.1}/PKG-INFO +5 -3
  2. clarifai-11.5.1/clarifai/__init__.py +1 -0
  3. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/app.py +2 -2
  4. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/__init__.py +2 -2
  5. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/dockerfile_template/Dockerfile.template +6 -2
  6. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/dummy_openai_model.py +59 -26
  7. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/mcp_class.py +14 -6
  8. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/model_builder.py +306 -10
  9. clarifai-11.5.1/clarifai/runners/models/openai_class.py +167 -0
  10. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/visual_classifier_class.py +1 -1
  11. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/code_script.py +4 -2
  12. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/const.py +3 -3
  13. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/data_utils.py +7 -1
  14. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/workflows/export.py +2 -5
  15. {clarifai-11.4.10 → clarifai-11.5.1/clarifai.egg-info}/PKG-INFO +5 -3
  16. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai.egg-info/requires.txt +4 -2
  17. {clarifai-11.4.10 → clarifai-11.5.1}/requirements.txt +4 -2
  18. clarifai-11.4.10/clarifai/__init__.py +0 -1
  19. clarifai-11.4.10/clarifai/runners/models/openai_class.py +0 -221
  20. {clarifai-11.4.10 → clarifai-11.5.1}/LICENSE +0 -0
  21. {clarifai-11.4.10 → clarifai-11.5.1}/MANIFEST.in +0 -0
  22. {clarifai-11.4.10 → clarifai-11.5.1}/README.md +0 -0
  23. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/README.md +0 -0
  24. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/__init__.py +0 -0
  25. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/__main__.py +0 -0
  26. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/base.py +0 -0
  27. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/compute_cluster.py +0 -0
  28. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/deployment.py +0 -0
  29. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/model.py +0 -0
  30. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/model_templates.py +0 -0
  31. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli/nodepool.py +0 -0
  32. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/cli.py +0 -0
  33. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/__init__.py +0 -0
  34. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/auth/__init__.py +0 -0
  35. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/auth/helper.py +0 -0
  36. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/auth/register.py +0 -0
  37. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/auth/stub.py +0 -0
  38. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/base.py +0 -0
  39. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/compute_cluster.py +0 -0
  40. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/dataset.py +0 -0
  41. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/deployment.py +0 -0
  42. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/input.py +0 -0
  43. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/lister.py +0 -0
  44. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/model.py +0 -0
  45. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/model_client.py +0 -0
  46. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/module.py +0 -0
  47. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/nodepool.py +0 -0
  48. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/runner.py +0 -0
  49. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/search.py +0 -0
  50. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/user.py +0 -0
  51. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/client/workflow.py +0 -0
  52. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/constants/base.py +0 -0
  53. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/constants/dataset.py +0 -0
  54. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/constants/input.py +0 -0
  55. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/constants/model.py +0 -0
  56. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/constants/rag.py +0 -0
  57. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/constants/search.py +0 -0
  58. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/constants/workflow.py +0 -0
  59. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/__init__.py +0 -0
  60. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/export/__init__.py +0 -0
  61. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/export/inputs_annotations.py +0 -0
  62. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/__init__.py +0 -0
  63. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/base.py +0 -0
  64. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/features.py +0 -0
  65. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/image.py +0 -0
  66. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/loaders/README.md +0 -0
  67. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/loaders/__init__.py +0 -0
  68. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/loaders/coco_captions.py +0 -0
  69. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/loaders/coco_detection.py +0 -0
  70. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/loaders/imagenet_classification.py +0 -0
  71. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/loaders/xview_detection.py +0 -0
  72. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/multimodal.py +0 -0
  73. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/text.py +0 -0
  74. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/datasets/upload/utils.py +0 -0
  75. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/errors.py +0 -0
  76. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/models/__init__.py +0 -0
  77. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/models/api.py +0 -0
  78. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/modules/README.md +0 -0
  79. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/modules/__init__.py +0 -0
  80. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/modules/css.py +0 -0
  81. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/modules/pages.py +0 -0
  82. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/modules/style.css +0 -0
  83. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/rag/__init__.py +0 -0
  84. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/rag/rag.py +0 -0
  85. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/rag/utils.py +0 -0
  86. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/__init__.py +0 -0
  87. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/model_class.py +0 -0
  88. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/model_run_locally.py +0 -0
  89. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/model_runner.py +0 -0
  90. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/model_servicer.py +0 -0
  91. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/models/visual_detector_class.py +0 -0
  92. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/server.py +0 -0
  93. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/__init__.py +0 -0
  94. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/data_types/__init__.py +0 -0
  95. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/data_types/data_types.py +0 -0
  96. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/loader.py +0 -0
  97. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/method_signatures.py +0 -0
  98. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/openai_convertor.py +0 -0
  99. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/serializers.py +0 -0
  100. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/runners/utils/url_fetcher.py +0 -0
  101. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/schema/search.py +0 -0
  102. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/urls/helper.py +0 -0
  103. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/__init__.py +0 -0
  104. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/cli.py +0 -0
  105. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/config.py +0 -0
  106. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/constants.py +0 -0
  107. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/evaluation/__init__.py +0 -0
  108. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/evaluation/helpers.py +0 -0
  109. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/evaluation/main.py +0 -0
  110. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/evaluation/testset_annotation_parser.py +0 -0
  111. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/logging.py +0 -0
  112. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/misc.py +0 -0
  113. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/model_train.py +0 -0
  114. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/utils/protobuf.py +0 -0
  115. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/versions.py +0 -0
  116. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/workflows/__init__.py +0 -0
  117. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/workflows/utils.py +0 -0
  118. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai/workflows/validate.py +0 -0
  119. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai.egg-info/SOURCES.txt +0 -0
  120. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai.egg-info/dependency_links.txt +0 -0
  121. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai.egg-info/entry_points.txt +0 -0
  122. {clarifai-11.4.10 → clarifai-11.5.1}/clarifai.egg-info/top_level.txt +0 -0
  123. {clarifai-11.4.10 → clarifai-11.5.1}/pyproject.toml +0 -0
  124. {clarifai-11.4.10 → clarifai-11.5.1}/setup.cfg +0 -0
  125. {clarifai-11.4.10 → clarifai-11.5.1}/setup.py +0 -0
  126. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_app.py +0 -0
  127. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_auth.py +0 -0
  128. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_data_upload.py +0 -0
  129. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_eval.py +0 -0
  130. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_misc.py +0 -0
  131. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_model_predict.py +0 -0
  132. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_model_train.py +0 -0
  133. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_modules.py +0 -0
  134. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_rag.py +0 -0
  135. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_search.py +0 -0
  136. {clarifai-11.4.10 → clarifai-11.5.1}/tests/test_stub.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: clarifai
3
- Version: 11.4.10
3
+ Version: 11.5.1
4
4
  Home-page: https://github.com/Clarifai/clarifai-python
5
5
  Author: Clarifai
6
6
  Author-email: support@clarifai.com
@@ -19,8 +19,8 @@ Classifier: Operating System :: OS Independent
19
19
  Requires-Python: >=3.8
20
20
  Description-Content-Type: text/markdown
21
21
  License-File: LICENSE
22
- Requires-Dist: clarifai-grpc>=11.3.4
23
- Requires-Dist: clarifai-protocol>=0.0.23
22
+ Requires-Dist: clarifai-grpc>=11.5.5
23
+ Requires-Dist: clarifai-protocol>=0.0.24
24
24
  Requires-Dist: numpy>=1.22.0
25
25
  Requires-Dist: tqdm>=4.65.0
26
26
  Requires-Dist: PyYAML>=6.0.1
@@ -31,6 +31,8 @@ Requires-Dist: fsspec>=2024.6.1
31
31
  Requires-Dist: click>=8.1.7
32
32
  Requires-Dist: requests>=2.32.3
33
33
  Requires-Dist: aiohttp>=3.10.0
34
+ Requires-Dist: uv==0.7.12
35
+ Requires-Dist: ruff==0.11.4
34
36
  Provides-Extra: all
35
37
  Requires-Dist: pycocotools>=2.0.7; extra == "all"
36
38
  Dynamic: author
@@ -0,0 +1 @@
1
+ __version__ = "11.5.1"
@@ -470,8 +470,8 @@ class App(Lister, BaseClient):
470
470
  model = self.model(
471
471
  model_id=node['model']['model_id'],
472
472
  model_version={"id": node['model'].get('model_version_id', "")},
473
- user_id=node['model'].get('user_id', ""),
474
- app_id=node['model'].get('app_id', ""),
473
+ user_id=node['model'].get('user_id', self.user_app_id.user_id),
474
+ app_id=node['model'].get('app_id', self.user_app_id.app_id),
475
475
  )
476
476
  except Exception as e:
477
477
  if "Model does not exist" in str(e):
@@ -1,11 +1,11 @@
1
- from .models.model_builder import ModelBuilder
1
+ from .models.mcp_class import MCPModelClass
2
2
  from .models.model_class import ModelClass
3
3
  from .models.model_runner import ModelRunner
4
4
  from .models.openai_class import OpenAIModelClass
5
5
 
6
6
  __all__ = [
7
7
  "ModelRunner",
8
- "ModelBuilder",
9
8
  "ModelClass",
9
+ "MCPModelClass",
10
10
  "OpenAIModelClass",
11
11
  ]
@@ -3,9 +3,13 @@ FROM --platform=$TARGETPLATFORM ${FINAL_IMAGE} as final
3
3
 
4
4
  COPY --link requirements.txt /home/nonroot/requirements.txt
5
5
 
6
+ ENV VIRTUAL_ENV=/venv
7
+ ENV PATH="/home/nonroot/.local/bin:$VIRTUAL_ENV/bin:$PATH"
8
+
9
+
6
10
  # Update clarifai package so we always have latest protocol to the API. Everything should land in /venv
7
- RUN ["pip", "install", "--no-cache-dir", "-r", "/home/nonroot/requirements.txt"]
8
- RUN ["pip", "show", "clarifai"]
11
+ RUN ["uv", "pip", "install", "--no-cache-dir", "-r", "/home/nonroot/requirements.txt"]
12
+ RUN ["uv", "pip", "show", "--no-cache-dir", "clarifai"]
9
13
 
10
14
  # Set the NUMBA cache dir to /tmp
11
15
  # Set the TORCHINDUCTOR cache dir to /tmp
@@ -13,9 +13,9 @@ class MockOpenAIClient:
13
13
  def create(self, **kwargs):
14
14
  """Mock create method for compatibility."""
15
15
  if kwargs.get("stream", False):
16
- return MockCompletionStream(kwargs.get("messages", []))
16
+ return MockCompletionStream(**kwargs)
17
17
  else:
18
- return MockCompletion(kwargs.get("messages", []))
18
+ return MockCompletion(**kwargs)
19
19
 
20
20
  def __init__(self):
21
21
  self.chat = self # Make self.chat point to self for compatibility
@@ -25,6 +25,19 @@ class MockOpenAIClient:
25
25
  class MockCompletion:
26
26
  """Mock completion object that mimics the OpenAI completion response structure."""
27
27
 
28
+ class Usage:
29
+ def __init__(self, prompt_tokens, completion_tokens, total_tokens):
30
+ self.total_tokens = total_tokens
31
+ self.prompt_tokens = prompt_tokens
32
+ self.completion_tokens = completion_tokens
33
+
34
+ def to_dict(self):
35
+ return dict(
36
+ total_tokens=self.total_tokens,
37
+ prompt_tokens=self.prompt_tokens,
38
+ completion_tokens=self.completion_tokens,
39
+ )
40
+
28
41
  class Choice:
29
42
  class Message:
30
43
  def __init__(self, content):
@@ -36,17 +49,21 @@ class MockCompletion:
36
49
  self.finish_reason = "stop"
37
50
  self.index = 0
38
51
 
39
- def __init__(self, messages):
52
+ def __init__(self, **kwargs):
40
53
  # Generate a simple response based on the last message
54
+ messages = kwargs.get("messages")
41
55
  last_message = messages[-1] if messages else {"content": ""}
42
56
  response_text = f"Echo: {last_message.get('content', '')}"
43
57
 
44
58
  self.choices = [self.Choice(response_text)]
45
- self.usage = {
46
- "prompt_tokens": len(str(messages)),
47
- "completion_tokens": len(response_text),
48
- "total_tokens": len(str(messages)) + len(response_text),
49
- }
59
+ self.usage = self.Usage(
60
+ **{
61
+ "prompt_tokens": len(str(messages)),
62
+ "completion_tokens": len(response_text),
63
+ "total_tokens": len(str(messages)) + len(response_text),
64
+ }
65
+ )
66
+
50
67
  self.id = "dummy-completion-id"
51
68
  self.created = 1234567890
52
69
  self.model = "dummy-model"
@@ -65,9 +82,12 @@ class MockCompletion:
65
82
  }
66
83
  for choice in self.choices
67
84
  ],
68
- "usage": self.usage,
85
+ "usage": self.usage.to_dict(),
69
86
  }
70
87
 
88
+ def model_dump(self):
89
+ return self.to_dict()
90
+
71
91
 
72
92
  class MockCompletionStream:
73
93
  """Mock completion stream that mimics the OpenAI streaming response structure."""
@@ -79,14 +99,27 @@ class MockCompletionStream:
79
99
  self.content = content
80
100
  self.role = "assistant" if content is None else None
81
101
 
102
+ class Usage:
103
+ def __init__(self, prompt_tokens, completion_tokens, total_tokens):
104
+ self.total_tokens = total_tokens
105
+ self.prompt_tokens = prompt_tokens
106
+ self.completion_tokens = completion_tokens
107
+
108
+ def to_dict(self):
109
+ return dict(
110
+ total_tokens=self.total_tokens,
111
+ prompt_tokens=self.prompt_tokens,
112
+ completion_tokens=self.completion_tokens,
113
+ )
114
+
82
115
  def __init__(self, content=None, include_usage=False):
83
116
  self.delta = self.Delta(content)
84
117
  self.finish_reason = None if content else "stop"
85
118
  self.index = 0
86
119
  self.usage = (
87
- {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
120
+ self.Usage(**{"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15})
88
121
  if include_usage
89
- else None
122
+ else self.Usage(None, None, None)
90
123
  )
91
124
 
92
125
  def __init__(self, content=None, include_usage=False):
@@ -114,11 +147,16 @@ class MockCompletionStream:
114
147
  ],
115
148
  }
116
149
  if self.usage:
117
- result["usage"] = self.usage
150
+ result["usage"] = self.usage.to_dict()
118
151
  return result
119
152
 
120
- def __init__(self, messages):
153
+ def model_dump(self):
154
+ return self.to_dict()
155
+
156
+ def __init__(self, **kwargs):
121
157
  # Generate a simple response based on the last message
158
+ messages = kwargs.get("messages")
159
+
122
160
  last_message = messages[-1] if messages else {"content": ""}
123
161
  self.response_text = f"Echo: {last_message.get('content', '')}"
124
162
  # Create chunks that ensure the full text is included in the first chunk
@@ -127,7 +165,7 @@ class MockCompletionStream:
127
165
  "", # Final chunk is empty to indicate completion
128
166
  ]
129
167
  self.current_chunk = 0
130
- self.include_usage = False
168
+ self.include_usage = kwargs.get("stream_options", {}).get("include_usage")
131
169
 
132
170
  def __iter__(self):
133
171
  return self
@@ -150,18 +188,14 @@ class DummyOpenAIModel(OpenAIModelClass):
150
188
  def _process_request(self, **kwargs) -> Dict[str, Any]:
151
189
  """Process a request for non-streaming responses."""
152
190
  completion_args = self._create_completion_args(kwargs)
153
- return self.client.chat.completions.create(**completion_args).to_dict()
191
+ return self.client.chat.completions.create(**completion_args).model_dump()
154
192
 
155
193
  def _process_streaming_request(self, **kwargs) -> Iterator[Dict[str, Any]]:
156
194
  """Process a request for streaming responses."""
157
- completion_args = self._create_completion_args(kwargs, stream=True)
158
- completion_stream = self.client.chat.completions.create(**completion_args)
159
- completion_stream.include_usage = kwargs.get('stream_options', {}).get(
160
- 'include_usage', False
161
- )
195
+ completion_stream = self.client.chat.completions.create(**kwargs)
162
196
 
163
197
  for chunk in completion_stream:
164
- yield chunk.to_dict()
198
+ yield chunk.model_dump()
165
199
 
166
200
  # Override the method directly for testing
167
201
  @OpenAIModelClass.method
@@ -169,14 +203,13 @@ class DummyOpenAIModel(OpenAIModelClass):
169
203
  """Direct implementation for testing purposes."""
170
204
  try:
171
205
  request_data = json.loads(req)
172
- params = self._extract_request_params(request_data)
173
-
206
+ request_data = self._create_completion_args(request_data)
174
207
  # Validate messages
175
- if not params.get("messages"):
208
+ if not request_data.get("messages"):
176
209
  yield "Error: No messages provided"
177
210
  return
178
211
 
179
- for message in params["messages"]:
212
+ for message in request_data["messages"]:
180
213
  if (
181
214
  not isinstance(message, dict)
182
215
  or "role" not in message
@@ -185,7 +218,7 @@ class DummyOpenAIModel(OpenAIModelClass):
185
218
  yield "Error: Invalid message format"
186
219
  return
187
220
 
188
- for chunk in self._process_streaming_request(**params):
221
+ for chunk in self._process_streaming_request(**request_data):
189
222
  yield json.dumps(chunk)
190
223
  except Exception as e:
191
224
  yield f"Error: {str(e)}"
@@ -2,14 +2,13 @@
2
2
 
3
3
  import asyncio
4
4
  import json
5
- from typing import Any
6
-
7
- from fastmcp import Client, FastMCP # use fastmcp v2 not the built in mcp
8
- from mcp import types
9
- from mcp.shared.exceptions import McpError
5
+ from typing import TYPE_CHECKING, Any
10
6
 
11
7
  from clarifai.runners.models.model_class import ModelClass
12
8
 
9
+ if TYPE_CHECKING:
10
+ from fastmcp import FastMCP
11
+
13
12
 
14
13
  class MCPModelClass(ModelClass):
15
14
  """Base class for wrapping FastMCP servers as a model running in Clarfai. This handles
@@ -19,10 +18,17 @@ class MCPModelClass(ModelClass):
19
18
  """
20
19
 
21
20
  def load_model(self):
21
+ try:
22
+ from fastmcp import Client
23
+ except ImportError:
24
+ raise ImportError(
25
+ "fastmcp package is required to use MCP functionality. "
26
+ "Install it with: pip install fastmcp"
27
+ )
22
28
  # in memory transport provided in fastmcp v2 so we can easily use the client functions.
23
29
  self.client = Client(self.get_server())
24
30
 
25
- def get_server(self) -> FastMCP:
31
+ def get_server(self) -> 'FastMCP':
26
32
  """Required method for each subclass to implement to return the FastMCP server to use."""
27
33
  raise NotImplementedError("Subclasses must implement get_server() method")
28
34
 
@@ -32,6 +38,8 @@ class MCPModelClass(ModelClass):
32
38
  return it's response.
33
39
 
34
40
  """
41
+ from mcp import types
42
+ from mcp.shared.exceptions import McpError
35
43
 
36
44
  async def send_notification(client_message: types.ClientNotification) -> None:
37
45
  async with self.client: