amsdal_ml 0.1.0__tar.gz → 0.1.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 (67) hide show
  1. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.gitignore +5 -2
  2. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/PKG-INFO +1 -1
  3. amsdal_ml-0.1.1/amsdal_ml/__about__.py +1 -0
  4. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/agents/agent.py +31 -7
  5. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/agents/default_qa_agent.py +20 -12
  6. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/agents/promts/react_chat.prompt +10 -1
  7. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/agents/retriever_tool.py +25 -6
  8. amsdal_ml-0.1.1/amsdal_ml/fileio/base_loader.py +63 -0
  9. amsdal_ml-0.1.1/amsdal_ml/fileio/openai_loader.py +69 -0
  10. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/mcp_client/stdio_client.py +22 -2
  11. amsdal_ml-0.1.1/amsdal_ml/mcp_server/server_retriever_stdio.py +59 -0
  12. amsdal_ml-0.1.1/amsdal_ml/migrations/0000_initial.py +36 -0
  13. amsdal_ml-0.1.1/amsdal_ml/ml_models/models.py +87 -0
  14. amsdal_ml-0.1.1/amsdal_ml/ml_models/openai_model.py +371 -0
  15. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_retrievers/default_retriever.py +1 -1
  16. amsdal_ml-0.1.1/latest-changelogs.md +5 -0
  17. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/pyproject.toml +4 -0
  18. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/tests/agents_tests/test_fakes.py +43 -23
  19. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/uv.lock +698 -491
  20. amsdal_ml-0.1.0/amsdal_ml/__about__.py +0 -1
  21. amsdal_ml-0.1.0/amsdal_ml/mcp_server/server_retriever_stdio.py +0 -11
  22. amsdal_ml-0.1.0/amsdal_ml/ml_models/models.py +0 -50
  23. amsdal_ml-0.1.0/amsdal_ml/ml_models/openai_model.py +0 -171
  24. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.amsdal/.dependencies +0 -0
  25. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.amsdal/.environment +0 -0
  26. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.amsdal/.secrets +0 -0
  27. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.amsdal-cli +0 -0
  28. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.github/workflows/ci.yml +0 -0
  29. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.github/workflows/release.yml +0 -0
  30. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/.github/workflows/tag_check.yml +0 -0
  31. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/README.md +0 -0
  32. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/Third-Party Materials - AMSDAL Dependencies - License Notices.md +0 -0
  33. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/__init__.py +0 -0
  34. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/agents/__init__.py +0 -0
  35. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/agents/promts/__init__.py +0 -0
  36. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/app.py +0 -0
  37. {amsdal_ml-0.1.0/amsdal_ml/mcp_client → amsdal_ml-0.1.1/amsdal_ml/fileio}/__init__.py +0 -0
  38. {amsdal_ml-0.1.0/amsdal_ml/mcp_server → amsdal_ml-0.1.1/amsdal_ml/mcp_client}/__init__.py +0 -0
  39. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/mcp_client/base.py +0 -0
  40. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/mcp_client/http_client.py +0 -0
  41. {amsdal_ml-0.1.0/amsdal_ml/ml_ingesting → amsdal_ml-0.1.1/amsdal_ml/mcp_server}/__init__.py +0 -0
  42. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_config.py +0 -0
  43. {amsdal_ml-0.1.0/amsdal_ml/ml_models → amsdal_ml-0.1.1/amsdal_ml/ml_ingesting}/__init__.py +0 -0
  44. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_ingesting/default_ingesting.py +0 -0
  45. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_ingesting/embedding_data.py +0 -0
  46. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_ingesting/ingesting.py +0 -0
  47. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_ingesting/openai_ingesting.py +0 -0
  48. {amsdal_ml-0.1.0/amsdal_ml/ml_retrievers → amsdal_ml-0.1.1/amsdal_ml/ml_models}/__init__.py +0 -0
  49. {amsdal_ml-0.1.0/amsdal_ml/models → amsdal_ml-0.1.1/amsdal_ml/ml_retrievers}/__init__.py +0 -0
  50. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_retrievers/openai_retriever.py +0 -0
  51. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/ml_retrievers/retriever.py +0 -0
  52. {amsdal_ml-0.1.0/tests → amsdal_ml-0.1.1/amsdal_ml/models}/__init__.py +0 -0
  53. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/models/embedding_model.py +0 -0
  54. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/amsdal_ml/py.typed +0 -0
  55. /amsdal_ml-0.1.0/latest-changelogs.md → /amsdal_ml-0.1.1/change-logs.md +0 -0
  56. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/config.yml +0 -0
  57. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/license_check.py +0 -0
  58. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/scripts/release.sh +0 -0
  59. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/scripts/tag_check.sh +0 -0
  60. {amsdal_ml-0.1.0/tests/agents_tests → amsdal_ml-0.1.1/tests}/__init__.py +0 -0
  61. /amsdal_ml-0.1.0/change-logs.md → /amsdal_ml-0.1.1/tests/agents_tests/__init__.py +0 -0
  62. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/tests/agents_tests/test_arun.py +0 -0
  63. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/tests/agents_tests/test_astream.py +0 -0
  64. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/tests/agents_tests/test_astream_final_only.py +0 -0
  65. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/tests/agents_tests/test_tool_call_arguments_async.py +0 -0
  66. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/tests/conftest.py +0 -0
  67. {amsdal_ml-0.1.0 → amsdal_ml-0.1.1}/tests/test_openai_model.py +0 -0
@@ -29,5 +29,8 @@ Thumbs.db
29
29
  !.env.example
30
30
 
31
31
 
32
- transactions/
33
- migrations/
32
+ /transactions/
33
+ /migrations/
34
+ /models/
35
+ /fixtures/
36
+ /static/
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: amsdal_ml
3
- Version: 0.1.0
3
+ Version: 0.1.1
4
4
  Summary: amsdal_ml plugin for AMSDAL Framework
5
5
  Requires-Python: >=3.11
6
6
  Requires-Dist: aiohttp==3.12.15
@@ -0,0 +1 @@
1
+ __version__ = '0.1.1'
@@ -6,10 +6,13 @@ from collections.abc import AsyncIterator
6
6
  from collections.abc import Iterator
7
7
  from typing import Any
8
8
  from typing import Literal
9
+ from typing import Optional
9
10
 
10
11
  from pydantic import BaseModel
11
12
  from pydantic import Field
12
13
 
14
+ from amsdal_ml.fileio.base_loader import FileAttachment
15
+
13
16
 
14
17
  class AgentMessage(BaseModel):
15
18
  role: Literal["SYSTEM", "USER", "ASSISTANT"]
@@ -24,16 +27,37 @@ class AgentOutput(BaseModel):
24
27
 
25
28
  class Agent(ABC):
26
29
  @abstractmethod
27
- async def arun(self, user_query: str) -> AgentOutput: ...
28
- @abstractmethod
29
- async def astream(self, user_query: str) -> AsyncIterator[str]:
30
- """Yield streamed chunks for the given query."""
31
- raise NotImplementedError
30
+ async def arun(
31
+ self,
32
+ user_query: str,
33
+ *,
34
+ attachments: Optional[list[FileAttachment]] = None,
35
+ ) -> AgentOutput:
36
+ ...
32
37
 
33
- def run(self, user_query: str) -> AgentOutput:
38
+ @abstractmethod
39
+ async def astream(
40
+ self,
41
+ user_query: str,
42
+ *,
43
+ attachments: Optional[list[FileAttachment]] = None,
44
+ ) -> AsyncIterator[str]:
45
+ ...
46
+
47
+ def run(
48
+ self,
49
+ user_query: str,
50
+ *,
51
+ attachments: Optional[list[FileAttachment]] = None,
52
+ ) -> AgentOutput:
34
53
  msg = "This agent is async-only. Use arun()."
35
54
  raise NotImplementedError(msg)
36
55
 
37
- def stream(self, user_query: str) -> Iterator[str]:
56
+ def stream(
57
+ self,
58
+ user_query: str,
59
+ *,
60
+ attachments: Optional[list[FileAttachment]] = None,
61
+ ) -> Iterator[str]:
38
62
  msg = "This agent is async-only. Use astream()."
39
63
  raise NotImplementedError(msg)
@@ -13,6 +13,7 @@ from typing import no_type_check
13
13
  from amsdal_ml.agents.agent import Agent
14
14
  from amsdal_ml.agents.agent import AgentOutput
15
15
  from amsdal_ml.agents.promts import get_prompt
16
+ from amsdal_ml.fileio.base_loader import FileAttachment
16
17
  from amsdal_ml.mcp_client.base import ToolClient
17
18
  from amsdal_ml.mcp_client.base import ToolInfo
18
19
  from amsdal_ml.ml_models.models import MLModel
@@ -30,6 +31,9 @@ _FINAL_RE = re.compile(
30
31
  r"Final Answer:\s*(?P<answer>.+)",
31
32
  re.DOTALL | re.IGNORECASE,
32
33
  )
34
+ # ---------- constants ----------
35
+
36
+ _MAX_PARSE_RETRIES = 5
33
37
 
34
38
 
35
39
 
@@ -186,11 +190,11 @@ class DefaultQAAgent(Agent):
186
190
  return str(content)
187
191
 
188
192
  # ---------- core run ----------
189
- def run(self, _user_query: str) -> AgentOutput:
193
+ def run(self, user_query: str, *, attachments: list[FileAttachment] | None = None) -> AgentOutput:
190
194
  msg = "DefaultQAAgent is async-only for now. Use arun()."
191
195
  raise NotImplementedError(msg)
192
196
 
193
- async def _run_async(self, user_query: str) -> AgentOutput:
197
+ async def _run_async(self, user_query: str, *, attachments: list[FileAttachment] | None = None) -> AgentOutput:
194
198
  if self._tool_clients and not self._tools_index_built:
195
199
  await self._build_clients_index()
196
200
 
@@ -198,10 +202,12 @@ class DefaultQAAgent(Agent):
198
202
  used_tools: list[str] = []
199
203
  parse_retries = 0
200
204
 
205
+
201
206
  for _ in range(self.max_steps):
202
207
  prompt = self._react_text(user_query, scratch)
203
- out = await self.model.ainvoke(prompt)
204
-
208
+ out = await self.model.ainvoke(prompt, attachments=attachments)
209
+ print("Model output:", out) # noqa: T201
210
+ print('promt:', prompt) # noqa: T201
205
211
  m_final = _FINAL_RE.search(out or "")
206
212
  if m_final:
207
213
  return AgentOutput(
@@ -213,7 +219,7 @@ class DefaultQAAgent(Agent):
213
219
  m_tool = _TOOL_CALL_RE.search(out or "")
214
220
  if not m_tool:
215
221
  parse_retries += 1
216
- if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
222
+ if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= _MAX_PARSE_RETRIES:
217
223
  msg = (
218
224
  "Invalid ReAct output. Expected EXACT format (Final or Tool-call). "
219
225
  f"Got:\n{out}"
@@ -221,6 +227,7 @@ class DefaultQAAgent(Agent):
221
227
  raise ValueError(
222
228
  msg
223
229
  )
230
+
224
231
  scratch += (
225
232
  "\nThought: Previous output violated the strict format. "
226
233
  "Reply again using EXACTLY one of the two specified formats.\n"
@@ -237,7 +244,7 @@ class DefaultQAAgent(Agent):
237
244
  raise ValueError(msg)
238
245
  except Exception as e:
239
246
  parse_retries += 1
240
- if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
247
+ if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= _MAX_PARSE_RETRIES:
241
248
  msg = f"Invalid Action Input JSON: {raw_input!r} ({e})"
242
249
  raise ValueError(msg) from e
243
250
  scratch += (
@@ -250,6 +257,7 @@ class DefaultQAAgent(Agent):
250
257
 
251
258
  try:
252
259
  result = await tool.run(args, context=None, convert_result=True)
260
+ print("Similarity search result:", result) # noqa: T201
253
261
  except Exception as e:
254
262
  # unified error payload
255
263
  err = {
@@ -276,12 +284,12 @@ class DefaultQAAgent(Agent):
276
284
  return self._stopped_response(used_tools)
277
285
 
278
286
  # ---------- public APIs ----------
279
- async def arun(self, user_query: str) -> AgentOutput:
280
- return await self._run_async(user_query)
287
+ async def arun(self, user_query: str, *, attachments: list[FileAttachment] | None = None) -> AgentOutput:
288
+ return await self._run_async(user_query, attachments=attachments)
281
289
 
282
290
  # ---------- streaming ----------
283
291
  @no_type_check
284
- async def astream(self, user_query: str) -> AsyncIterator[str]:
292
+ async def astream(self, user_query: str, *, attachments: list[FileAttachment] | None = None) -> AsyncIterator[str]:
285
293
  if self._tool_clients and not self._tools_index_built:
286
294
  await self._build_clients_index()
287
295
 
@@ -296,7 +304,7 @@ class DefaultQAAgent(Agent):
296
304
 
297
305
  # Normalize model.astream: it might be an async iterator already,
298
306
  # or a coroutine (or nested coroutines) that resolves to one.
299
- _val = self.model.astream(prompt)
307
+ _val = self.model.astream(prompt, attachments=attachments)
300
308
  while inspect.iscoroutine(_val):
301
309
  _val = await _val
302
310
 
@@ -320,7 +328,7 @@ class DefaultQAAgent(Agent):
320
328
  m_tool = _TOOL_CALL_RE.search(buffer or "")
321
329
  if not m_tool:
322
330
  parse_retries += 1
323
- if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
331
+ if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= _MAX_PARSE_RETRIES:
324
332
  msg = f"Invalid ReAct output (stream). Expected EXACT format. Got:\n{buffer}"
325
333
  raise ValueError(msg)
326
334
  scratch += (
@@ -339,7 +347,7 @@ class DefaultQAAgent(Agent):
339
347
  raise ValueError(msg)
340
348
  except Exception as e:
341
349
  parse_retries += 1
342
- if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= 1:
350
+ if self.on_parse_error == ParseErrorMode.RAISE or parse_retries >= _MAX_PARSE_RETRIES:
343
351
  msg = f"Invalid Action Input JSON: {raw_input!r} ({e})"
344
352
  raise ValueError(msg) from e
345
353
  scratch += (
@@ -21,7 +21,16 @@ RULES
21
21
  - `Action Input` MUST be a valid ONE-LINE JSON object (e.g. {{"a": 1, "b": 2}}).
22
22
  - Do NOT add anything before/after the block.
23
23
  - Do NOT print "Observation". The system will add it after tool execution.
24
-
24
+ - You CAN read and use any attached files provided by the platform (including PDFs). Do not claim inability to view files.
25
+ - Prefer not to use tools when the answer is fully supported by the user’s message and/or attached files. Use a tool only if it materially improves accuracy or retrieves missing facts.
26
+ - If information is insufficient, answer concisely with: “I don’t have enough data to answer.”
27
+ - Be deterministic and concise. No preambles, no meta-commentary, no Markdown.
28
+ - Numbers: preserve exact figures, units, and percentages; do not round unless asked.
29
+ - If you choose a tool:
30
+ - Use exactly one tool per step.
31
+ - The Action Input MUST be a single-line valid JSON object that matches the tool’s schema.
32
+ - If a tool errors or returns nothing useful, switch to the other format and produce a Final Answer.
33
+
25
34
  PREVIOUS CONVERSATION
26
35
  --------------------
27
36
  {chat_history}
@@ -1,5 +1,7 @@
1
1
  from __future__ import annotations
2
2
 
3
+ import logging
4
+ import sys
3
5
  from typing import Any
4
6
  from typing import Optional
5
7
 
@@ -9,6 +11,14 @@ from pydantic import Field
9
11
 
10
12
  from amsdal_ml.ml_retrievers.openai_retriever import OpenAIRetriever
11
13
 
14
+ logging.basicConfig(
15
+ level=logging.INFO,
16
+ format='%(asctime)s [%(levelname)s] %(message)s',
17
+ handlers=[
18
+ logging.FileHandler("server2.log"),
19
+ logging.StreamHandler(sys.stdout)
20
+ ]
21
+ )
12
22
 
13
23
  class RetrieverArgs(BaseModel):
14
24
  query: str = Field(..., description='User search query')
@@ -20,13 +30,22 @@ class RetrieverArgs(BaseModel):
20
30
  _retriever = OpenAIRetriever()
21
31
 
22
32
 
23
- async def retriever_search(args: RetrieverArgs) -> list[dict[str, Any]]:
33
+ async def retriever_search(
34
+ query: str = Field(..., description='User search query'),
35
+ k: int = 5,
36
+ include_tags: Optional[list[str]] = None,
37
+ exclude_tags: Optional[list[str]] = None,
38
+ ) -> list[dict[str, Any]]:
39
+ logging.info(
40
+ f"retriever_search called with query={query}, k={k}, include_tags={include_tags}, exclude_tags={exclude_tags}"
41
+ )
24
42
  chunks = await _retriever.asimilarity_search(
25
- query=args.query,
26
- k=args.k,
27
- include_tags=args.include_tags,
28
- exclude_tags=args.exclude_tags,
43
+ query=query,
44
+ k=k,
45
+ include_tags=include_tags,
46
+ exclude_tags=exclude_tags,
29
47
  )
48
+ logging.info(f"retriever_search found {len(chunks)} chunks: {chunks}")
30
49
  out: list[dict[str, Any]] = []
31
50
  for c in chunks:
32
51
  if hasattr(c, 'model_dump'):
@@ -42,7 +61,7 @@ async def retriever_search(args: RetrieverArgs) -> list[dict[str, Any]]:
42
61
 
43
62
  retriever_tool = Tool.from_function(
44
63
  retriever_search,
45
- name='retriever.search',
64
+ name='search',
46
65
  description='Semantic search in knowledge base (OpenAI embeddings)',
47
66
  structured_output=True,
48
67
  )
@@ -0,0 +1,63 @@
1
+ from __future__ import annotations
2
+
3
+ from abc import ABC
4
+ from abc import abstractmethod
5
+ from dataclasses import dataclass
6
+ from typing import Any
7
+ from typing import BinaryIO
8
+ from typing import Optional
9
+
10
+ from pydantic import BaseModel
11
+
12
+ PLAIN_TEXT = "plain_text"
13
+ FILE_ID = "file_id"
14
+
15
+
16
+ class FileData(BaseModel):
17
+ name: str
18
+ size: int | None = None
19
+ mime: Optional[str] = None
20
+
21
+
22
+ @dataclass
23
+ class FileAttachment:
24
+ type: str # one of: PLAIN_TEXT, FILE_ID
25
+ content: Any
26
+ metadata: dict[str, Any] | None = None
27
+
28
+ def __post_init__(self) -> None:
29
+ if self.metadata is None:
30
+ self.metadata = {}
31
+
32
+
33
+ @dataclass
34
+ class FileItem:
35
+ file: BinaryIO
36
+ filename: str | None = None
37
+ filedata: FileData | None = None
38
+
39
+ @staticmethod
40
+ def from_path(path: str, *, filedata: FileData | None = None) -> FileItem:
41
+ # Caller is responsible for lifecycle; loaders may close after upload.
42
+ f = open(path, "rb")
43
+ return FileItem(file=f, filename=path.split("/")[-1], filedata=filedata)
44
+
45
+ @staticmethod
46
+ def from_bytes(data: bytes, *, filename: str | None = None, filedata: FileData | None = None) -> FileItem:
47
+ import io
48
+ return FileItem(file=io.BytesIO(data), filename=filename, filedata=filedata)
49
+
50
+ @staticmethod
51
+ def from_str(text: str, *, filename: str | None = None, filedata: FileData | None = None) -> FileItem:
52
+ import io
53
+ return FileItem(file=io.BytesIO(text.encode("utf-8")), filename=filename, filedata=filedata)
54
+
55
+
56
+ class BaseFileLoader(ABC):
57
+ @abstractmethod
58
+ async def load(self, item: FileItem) -> FileAttachment:
59
+ """Upload a single file and return an attachment reference."""
60
+
61
+ @abstractmethod
62
+ async def load_batch(self, items: list[FileItem]) -> list[FileAttachment]:
63
+ """Upload multiple files; input and output are lists for simplicity."""
@@ -0,0 +1,69 @@
1
+ from __future__ import annotations
2
+
3
+ import asyncio
4
+ import io
5
+ import logging
6
+ from collections.abc import Sequence
7
+ from typing import Any
8
+ from typing import BinaryIO
9
+ from typing import Literal
10
+
11
+ from openai import AsyncOpenAI
12
+
13
+ from amsdal_ml.fileio.base_loader import FILE_ID
14
+ from amsdal_ml.fileio.base_loader import BaseFileLoader
15
+ from amsdal_ml.fileio.base_loader import FileAttachment
16
+ from amsdal_ml.fileio.base_loader import FileData
17
+ from amsdal_ml.fileio.base_loader import FileItem
18
+
19
+ logger = logging.getLogger(__name__)
20
+
21
+ AllowedPurpose = Literal["assistants", "batch", "fine-tune", "vision", "user_data", "evals"]
22
+
23
+
24
+ class OpenAIFileLoader(BaseFileLoader):
25
+ """
26
+ Loader which uploads files into OpenAI Files API and returns openai_file_id.
27
+ """
28
+
29
+ def __init__(self, client: AsyncOpenAI, *, purpose: AllowedPurpose = "assistants") -> None:
30
+ self.client = client
31
+ self.purpose: AllowedPurpose = purpose # mypy: Literal union, matches SDK
32
+
33
+ async def _upload_one(self, file: BinaryIO, *, filename: str | None, filedata: FileData | None) -> FileAttachment:
34
+ try:
35
+ if hasattr(file, "seek"):
36
+ file.seek(0)
37
+ except Exception as exc: # pragma: no cover
38
+ logger.debug("seek(0) failed for %r: %s", filename or file, exc)
39
+
40
+ buf = file if isinstance(file, io.BytesIO) else io.BytesIO(file.read())
41
+
42
+ up = await self.client.files.create(file=(filename or "upload.bin", buf), purpose=self.purpose)
43
+
44
+ meta: dict[str, Any] = {
45
+ "filename": filename,
46
+ "provider": "openai",
47
+ "file": {
48
+ "id": up.id,
49
+ "bytes": getattr(up, "bytes", None),
50
+ "purpose": getattr(up, "purpose", self.purpose),
51
+ "created_at": getattr(up, "created_at", None),
52
+ "status": getattr(up, "status", None),
53
+ "status_details": getattr(up, "status_details", None),
54
+ },
55
+ }
56
+ if filedata is not None:
57
+ meta["filedata"] = filedata.model_dump()
58
+
59
+ return FileAttachment(type=FILE_ID, content=up.id, metadata=meta)
60
+
61
+ async def load(self, item: FileItem) -> FileAttachment:
62
+ return await self._upload_one(item.file, filename=item.filename, filedata=item.filedata)
63
+
64
+ async def load_batch(self, items: Sequence[FileItem]) -> list[FileAttachment]:
65
+ tasks = [
66
+ asyncio.create_task(self._upload_one(it.file, filename=it.filename, filedata=it.filedata))
67
+ for it in items
68
+ ]
69
+ return await asyncio.gather(*tasks)
@@ -1,10 +1,12 @@
1
1
  from __future__ import annotations
2
2
 
3
3
  import asyncio
4
+ import base64
4
5
  from collections.abc import Iterable
5
6
  from contextlib import AsyncExitStack
6
7
  from typing import Any
7
8
 
9
+ from amsdal_utils.config.manager import AmsdalConfigManager
8
10
  from mcp import ClientSession
9
11
  from mcp import StdioServerParameters
10
12
  from mcp.client.stdio import stdio_client
@@ -19,7 +21,14 @@ class StdioClient(ToolClient):
19
21
  MCP over STDIO client.
20
22
  """
21
23
 
22
- def __init__(self, alias: str, module_or_cmd: str, *args: str, persist_session: bool = True):
24
+ def __init__(
25
+ self,
26
+ alias: str,
27
+ module_or_cmd: str,
28
+ *args: str,
29
+ persist_session: bool = True,
30
+ send_amsdal_config: bool = True,
31
+ ):
23
32
  self.alias = alias
24
33
  if module_or_cmd in ("python", "python3"):
25
34
  self._command = module_or_cmd
@@ -27,7 +36,9 @@ class StdioClient(ToolClient):
27
36
  else:
28
37
  self._command = "python"
29
38
  self._args = ["-m", module_or_cmd]
30
-
39
+ if send_amsdal_config:
40
+ self._args.append("--amsdal-config")
41
+ self._args.append(self._build_amsdal_config_arg())
31
42
  self._persist = persist_session
32
43
  self._lock = asyncio.Lock()
33
44
  self._stack: AsyncExitStack | None = None
@@ -115,6 +126,7 @@ class StdioClient(ToolClient):
115
126
  rx, tx = await stack.enter_async_context(stdio_client(params))
116
127
  s = await stack.enter_async_context(ClientSession(rx, tx))
117
128
  await s.initialize()
129
+ print("Calling tool:", tool_name, "with args:", args) # noqa: T201
118
130
  res = await self._call_with_timeout(s.call_tool(tool_name, args), timeout=timeout)
119
131
  return getattr(res, "content", res)
120
132
 
@@ -128,3 +140,11 @@ class StdioClient(ToolClient):
128
140
  s = await self._ensure_session()
129
141
  res = await self._call_with_timeout(s.call_tool(tool_name, args), timeout=timeout)
130
142
  return getattr(res, "content", res)
143
+
144
+ def _build_amsdal_config_arg(self) -> str:
145
+ """
146
+ Build a JSON string argument representing the current Amsdal configuration.
147
+ This can be passed to the subprocess to ensure it has the same configuration context.
148
+ """
149
+ config = AmsdalConfigManager().get_config()
150
+ return base64.b64encode(config.model_dump_json().encode('utf-8')).decode('utf-8')
@@ -0,0 +1,59 @@
1
+ import argparse
2
+ import asyncio
3
+ import base64
4
+ import json
5
+ import logging
6
+ import sys
7
+ from typing import Any
8
+ from typing import cast
9
+
10
+ from amsdal.manager import AmsdalManager
11
+ from amsdal.manager import AsyncAmsdalManager
12
+ from amsdal_utils.config.data_models.amsdal_config import AmsdalConfig
13
+ from amsdal_utils.config.manager import AmsdalConfigManager
14
+ from mcp.server.fastmcp import FastMCP
15
+
16
+ from amsdal_ml.agents.retriever_tool import retriever_search
17
+
18
+ logging.basicConfig(
19
+ level=logging.INFO,
20
+ format='%(asctime)s [%(levelname)s] %(message)s',
21
+ handlers=[
22
+ logging.FileHandler("server.log"),
23
+ logging.StreamHandler(sys.stdout)
24
+ ]
25
+ )
26
+
27
+ parser = argparse.ArgumentParser()
28
+ parser.add_argument("--amsdal-config", required=False, help="Base64-encoded config string")
29
+ args = parser.parse_args()
30
+
31
+ logging.info(f"Starting server with args: {args}")
32
+
33
+ if args.amsdal_config:
34
+ decoded = base64.b64decode(args.amsdal_config).decode("utf-8")
35
+ amsdal_config = AmsdalConfig(**json.loads(decoded))
36
+ logging.info(f"Loaded Amsdal config: {amsdal_config}")
37
+ AmsdalConfigManager().set_config(amsdal_config)
38
+
39
+ manager: Any
40
+ if amsdal_config.async_mode:
41
+ manager = AsyncAmsdalManager()
42
+ logging.info("pre-setup")
43
+ asyncio.run(cast(Any, manager).setup())
44
+ logging.info("post-setup")
45
+ asyncio.run(cast(Any, manager).post_setup())
46
+ logging.info("manager inited")
47
+ else:
48
+ manager = AmsdalManager()
49
+ cast(Any, manager).setup()
50
+ cast(Any, manager).post_setup()
51
+
52
+ server = FastMCP('retriever-stdio')
53
+ server.tool(
54
+ name='search',
55
+ description='Semantic search in knowledge base (OpenAI embeddings)',
56
+ structured_output=True,
57
+ )(retriever_search)
58
+
59
+ server.run(transport='stdio')
@@ -0,0 +1,36 @@
1
+ from amsdal_models.migration import migrations
2
+ from amsdal_utils.models.enums import ModuleType
3
+
4
+
5
+ class Migration(migrations.Migration):
6
+ operations: list[migrations.Operation] = [
7
+ migrations.CreateClass(
8
+ module_type=ModuleType.CONTRIB,
9
+ class_name="EmbeddingModel",
10
+ new_schema={
11
+ "title": "EmbeddingModel",
12
+ "required": ["object_class", "object_id", "chunk_index", "raw_text", "embedding"],
13
+ "properties": {
14
+ "object_class": {"type": "string", "title": "Linked object class"},
15
+ "object_id": {"type": "string", "title": "Linked object ID"},
16
+ "chunk_index": {"type": "integer", "title": "Chunk index"},
17
+ "raw_text": {"type": "string", "title": "Raw text used for embedding"},
18
+ "embedding": {
19
+ "type": "array",
20
+ "items": {"type": "number"},
21
+ "title": "Embedding",
22
+ "additional_type": "vector",
23
+ "dimensions": 1536,
24
+ },
25
+ "tags": {"type": "array", "items": {"type": "string"}, "title": "Embedding tags"},
26
+ "ml_metadata": {"type": "anything", "title": "ML metadata"},
27
+ },
28
+ "storage_metadata": {
29
+ "table_name": "embedding_model",
30
+ "db_fields": {},
31
+ "primary_key": ["partition_key"],
32
+ "foreign_keys": {},
33
+ },
34
+ },
35
+ ),
36
+ ]
@@ -0,0 +1,87 @@
1
+ from __future__ import annotations
2
+
3
+ from abc import ABC
4
+ from abc import abstractmethod
5
+ from collections.abc import AsyncIterator
6
+
7
+ from amsdal_ml.fileio.base_loader import PLAIN_TEXT
8
+ from amsdal_ml.fileio.base_loader import FileAttachment
9
+
10
+
11
+ class ModelError(Exception):
12
+ """Base exception for all ML models."""
13
+
14
+
15
+ class ModelConnectionError(ModelError):
16
+ """Network or connection failure to the provider (timeouts, DNS, TLS, etc.)."""
17
+
18
+
19
+ class ModelRateLimitError(ModelError):
20
+ """Provider's rate limit reached (HTTP 429)."""
21
+
22
+
23
+ class ModelAPIError(ModelError):
24
+ """API responded with an error (any 4xx/5xx, except 429)."""
25
+
26
+
27
+ class MLModel(ABC):
28
+ @abstractmethod
29
+ def setup(self) -> None:
30
+ """Initialize any clients or resources needed before inference."""
31
+ raise NotImplementedError
32
+
33
+ @abstractmethod
34
+ def teardown(self) -> None:
35
+ """Clean up resources after use."""
36
+ raise NotImplementedError
37
+
38
+ def supported_attachments(self) -> set[str]:
39
+ """Return a set of universal attachment kinds, e.g. {PLAIN_TEXT, FILE_ID}."""
40
+ return {PLAIN_TEXT}
41
+
42
+ @abstractmethod
43
+ def invoke(
44
+ self,
45
+ prompt: str,
46
+ *,
47
+ attachments: list[FileAttachment] | None = None,
48
+ ) -> str:
49
+ """Run synchronous inference with the model."""
50
+ raise NotImplementedError
51
+
52
+ @abstractmethod
53
+ async def ainvoke(
54
+ self,
55
+ prompt: str,
56
+ *,
57
+ attachments: list[FileAttachment] | None = None,
58
+ ) -> str:
59
+ """Run asynchronous inference with the model."""
60
+ raise NotImplementedError
61
+
62
+ @abstractmethod
63
+ def stream(
64
+ self,
65
+ prompt: str,
66
+ *,
67
+ attachments: list[FileAttachment] | None = None,
68
+ ):
69
+ """Stream synchronous inference results from the model."""
70
+ raise NotImplementedError
71
+
72
+ @abstractmethod
73
+ def astream(
74
+ self,
75
+ prompt: str,
76
+ *,
77
+ attachments: list[FileAttachment] | None = None,
78
+ ) -> AsyncIterator[str]:
79
+ """
80
+ Stream asynchronous inference results as an async generator.
81
+
82
+ Subclasses should implement this like:
83
+
84
+ async def astream(... ) -> AsyncIterator[str]:
85
+ yield "chunk"
86
+ """
87
+ raise NotImplementedError