goodmem-autogen 0.3.0__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.
@@ -0,0 +1,216 @@
1
+ Metadata-Version: 2.4
2
+ Name: goodmem-autogen
3
+ Version: 0.3.0
4
+ Summary: GoodMem memory and tools for the AutoGen agent framework.
5
+ Keywords: autogen,goodmem,memory,rag,agents,llm
6
+ Author: PAIR Systems
7
+ Requires-Python: >=3.10
8
+ Description-Content-Type: text/markdown
9
+ Classifier: Programming Language :: Python :: 3
10
+ Classifier: Programming Language :: Python :: 3.10
11
+ Classifier: Programming Language :: Python :: 3.11
12
+ Classifier: Programming Language :: Python :: 3.12
13
+ Classifier: License :: OSI Approved :: MIT License
14
+ Classifier: Operating System :: OS Independent
15
+ Classifier: Topic :: Software Development :: Libraries :: Python Modules
16
+ Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
17
+ Requires-Dist: autogen-core>=0.7.5
18
+ Requires-Dist: goodmem>=0.1.34
19
+ Requires-Dist: pydantic>=2.0,<3
20
+ Requires-Dist: typing-extensions>=4.7
21
+ Requires-Dist: pytest>=7 ; extra == "dev"
22
+ Requires-Dist: pytest-asyncio>=0.23 ; extra == "dev"
23
+ Requires-Dist: pytest-timeout ; extra == "dev"
24
+ Requires-Dist: httpx ; extra == "dev"
25
+ Requires-Dist: ruff>=0.6 ; extra == "dev"
26
+ Requires-Dist: mypy>=1.11 ; extra == "dev"
27
+ Requires-Dist: build ; extra == "dev"
28
+ Requires-Dist: twine ; extra == "dev"
29
+ Project-URL: Homepage, https://github.com/PAIR-Systems-Inc/goodmem-autogen
30
+ Project-URL: Issues, https://github.com/PAIR-Systems-Inc/goodmem-autogen/issues
31
+ Project-URL: Repository, https://github.com/PAIR-Systems-Inc/goodmem-autogen
32
+ Provides-Extra: dev
33
+
34
+ # goodmem-autogen
35
+
36
+ [GoodMem](https://goodmem.ai) memory and tools for the
37
+ [AutoGen](https://github.com/microsoft/autogen) agent framework.
38
+
39
+ Two ways in:
40
+
41
+ 1. **`GoodMemContextProvider`** — an `autogen_core.memory.Memory` backed by a
42
+ GoodMem space, so relevant passages are injected into the model context on
43
+ every turn.
44
+ 2. **`create_goodmem_search_tool` / `create_goodmem_admin_tools`** — function
45
+ tools an agent can call directly.
46
+
47
+ Built on the official `goodmem` SDK's async client, so nothing blocks the
48
+ event loop.
49
+
50
+ ## Install
51
+
52
+ ```bash
53
+ pip install goodmem-autogen
54
+ ```
55
+
56
+ Requires Python 3.10+, `autogen-core` 0.7.5+ and the `goodmem` SDK 0.1.34+
57
+ (installed with it). Version 0.2 is a break from
58
+ 0.1 — see [CHANGELOG](CHANGELOG.md) for the mapping.
59
+
60
+ ## As an AutoGen `Memory`
61
+
62
+ ```python
63
+ from autogen_core.memory import MemoryContent, MemoryMimeType
64
+ from goodmem_autogen import GoodMemContextProvider, GoodMemMemoryConfig
65
+
66
+ provider = GoodMemContextProvider(
67
+ config=GoodMemMemoryConfig(
68
+ base_url="https://goodmem.example.com",
69
+ api_key="gm_...", # stored as SecretStr, never serialized
70
+ space_name="handbook", # or space_id="..." to skip the lookup
71
+ embedder_id="<embedder-uuid>",
72
+ )
73
+ )
74
+
75
+ await provider.add(MemoryContent(
76
+ content="Refunds above $500 need a manager's approval.",
77
+ mime_type=MemoryMimeType.TEXT,
78
+ metadata={"title": "handbook", "category": "policy"},
79
+ ))
80
+
81
+ results = await provider.query("who approves a large refund?")
82
+ await provider.close()
83
+ ```
84
+
85
+ `add()` waits for the memory to finish indexing by default, so a query right
86
+ after it finds the result. Searching is never used as a way to wait.
87
+
88
+ Attach it to an agent and `update_context` injects retrieved passages as a
89
+ system message each turn — the same pattern AutoGen's own `ListMemory` uses.
90
+
91
+ ### What a result carries
92
+
93
+ `MemoryContent` has no score field, so provenance lives in `metadata`:
94
+
95
+ ```python
96
+ {
97
+ "title": "handbook", "category": "policy", # the memory's own metadata
98
+ "chunk_id": "...", "memory_id": "...", "space_id": "...", "source": "...",
99
+ "score": -0.53, "score_kind": "vector",
100
+ "partial": False, "statuses": [],
101
+ }
102
+ ```
103
+
104
+ - `partial` is `True` when part of the search did not complete — a reranker
105
+ was unavailable, one space was unreachable. The passages are usable but may
106
+ be incomplete, and `statuses` says why. A search that produced nothing
107
+ usable returns an empty result and emits a warning carrying the statuses,
108
+ so it is distinguishable from "no matches" without being raised. The
109
+ search tool returns `partial: true` with `statuses` in its JSON for the
110
+ same case.
111
+ - `score` is passed through exactly as GoodMem reports it. Vector scores are
112
+ opaque similarities that may be negative; reranker scores are on a scale
113
+ that depends on the reranker model (Voyage `rerank-2.5` ~`0.27..0.93`, Jina
114
+ `jina-reranker-v3` ~`-0.14..0.43` on the same documents). `score_kind` says
115
+ which you have — which is why `relevance_threshold` requires a
116
+ `reranker_id`, and why it must be calibrated for the reranker in use rather
117
+ than assumed to be 0–1. The threshold is applied by the server; if it
118
+ removes every result the provider warns, since an empty result would
119
+ otherwise read as "no matches".
120
+ - `score_kind` is read from the response, not the configuration. If the
121
+ reranker fails (`RERANKING_FAILED`, or `NOT_FOUND` for the reranker) the
122
+ server still returns the vector-search hits; they are kept, labelled
123
+ `vector`, and marked `partial` with those statuses.
124
+
125
+ ## As tools
126
+
127
+ The tools take a `goodmem.AsyncGoodmem` client, which you create and close.
128
+
129
+ ```python
130
+ from autogen_core import CancellationToken
131
+ from goodmem_autogen import create_goodmem_admin_tools, create_goodmem_search_tool
132
+ from goodmem import AsyncGoodmem
133
+
134
+ client = AsyncGoodmem(
135
+ base_url="https://goodmem.example.com",
136
+ api_key="gm_...",
137
+ # verify="/path/to/ca.pem", # a server with a self-signed certificate
138
+ )
139
+
140
+ search = create_goodmem_search_tool(
141
+ client, space_ids=["<space-uuid>"], limit=5,
142
+ reranker_id="<reranker-uuid>", # optional
143
+ metadata_filter={"category": "policy"}, # optional, escaped for you
144
+ )
145
+ admin_tools = create_goodmem_admin_tools(client) # only for agents that need them
146
+
147
+ # What an agent's tool call does:
148
+ print(await search.run_json({"query": "who approves a large refund?"}, CancellationToken()))
149
+
150
+ await client.close()
151
+ ```
152
+
153
+ The model supplies only the query; spaces, reranking and filters are yours, so
154
+ an agent cannot redirect a search or widen it mid-run. The search tool returns
155
+ JSON: `results` (each with `chunk_text`, `score`, `score_kind`, `metadata` and
156
+ IDs), `total_results`, `partial`, and `statuses` when the server reported any.
157
+
158
+ `create_goodmem_admin_tools(client)` adds space and memory management
159
+ (`goodmem_list_embedders`, `goodmem_list_rerankers`, `goodmem_list_spaces`,
160
+ `goodmem_get_space`, `goodmem_create_space`, `goodmem_update_space`,
161
+ `goodmem_delete_space`, `goodmem_create_memory`, `goodmem_list_memories`,
162
+ `goodmem_get_memory`, `goodmem_delete_memory`, and `goodmem_upload_file`
163
+ when `upload_dir` is given). A listing returns at most `max_items` (default
164
+ 100) and sets `truncated` when that cap may have cut it short. These
165
+ carry the authority of the configured API key — give them only to agents that
166
+ need them. File upload is only created when you pass `upload_dir`, and paths
167
+ resolving outside that directory are refused before the file is opened.
168
+
169
+ Every ID — a tool argument, `space_ids`, `reranker_id`, or an ID field of the
170
+ config — must be a UUID (it is lowercased); anything else raises `ValueError`
171
+ before a request is made. The SDK puts IDs into request paths unescaped, so
172
+ `"../spaces/<id>"` given as a memory ID would otherwise address a whole space.
173
+
174
+ ## Cancellation
175
+
176
+ `add`, `add_file` and `query` honour an `autogen_core.CancellationToken`: an
177
+ already-cancelled token prevents the request, and cancelling mid-flight aborts
178
+ it.
179
+
180
+ ## Clearing a space
181
+
182
+ `clear()` deletes every memory in the space and requires
183
+ `allow_clear=True` on the config, so a reflexive `clear()` cannot empty a
184
+ space by accident.
185
+
186
+ ## Filters
187
+
188
+ A filter is a GoodMem expression applied to every configured space, e.g.
189
+ `CAST(val('$.category') AS TEXT) = 'policy'`. Pass `metadata_filter={...}` and
190
+ it is built and escaped for you. Writing one by hand: inside a quoted value
191
+ escape `'` as `\'` and `\` as `\\` — SQL-style `''` doubling is rejected by
192
+ the server.
193
+
194
+ ## Development
195
+
196
+ ```bash
197
+ uv venv && uv pip install -e ".[dev]"
198
+ uv run ruff check goodmem_autogen tests
199
+ uv run mypy goodmem_autogen
200
+ uv run pytest -m "not integration" # offline: the real SDK over a mock transport or a local server
201
+ GOODMEM_BASE_URL=... GOODMEM_API_KEY=... GOODMEM_EMBEDDER_ID=... \
202
+ GOODMEM_VERIFY_SSL=false uv run pytest -m integration
203
+ ```
204
+
205
+ CI runs the first four on Python 3.10–3.13, then `uv build` and an import
206
+ of the wheel in a clean environment. It does not run the live suite, which
207
+ needs a server: `GOODMEM_EMBEDDER_ID` is the embedder its spaces are
208
+ created with, and `GOODMEM_VERIFY_SSL=false` is for a local server with a
209
+ self-signed certificate. The offline suite also executes every Python
210
+ snippet in this README against a local server.
211
+
212
+ Offline tests use event shapes captured from a live server. There is no
213
+ default API key — live tests skip unless the environment provides one.
214
+
215
+ MIT.
216
+
@@ -0,0 +1,182 @@
1
+ # goodmem-autogen
2
+
3
+ [GoodMem](https://goodmem.ai) memory and tools for the
4
+ [AutoGen](https://github.com/microsoft/autogen) agent framework.
5
+
6
+ Two ways in:
7
+
8
+ 1. **`GoodMemContextProvider`** — an `autogen_core.memory.Memory` backed by a
9
+ GoodMem space, so relevant passages are injected into the model context on
10
+ every turn.
11
+ 2. **`create_goodmem_search_tool` / `create_goodmem_admin_tools`** — function
12
+ tools an agent can call directly.
13
+
14
+ Built on the official `goodmem` SDK's async client, so nothing blocks the
15
+ event loop.
16
+
17
+ ## Install
18
+
19
+ ```bash
20
+ pip install goodmem-autogen
21
+ ```
22
+
23
+ Requires Python 3.10+, `autogen-core` 0.7.5+ and the `goodmem` SDK 0.1.34+
24
+ (installed with it). Version 0.2 is a break from
25
+ 0.1 — see [CHANGELOG](CHANGELOG.md) for the mapping.
26
+
27
+ ## As an AutoGen `Memory`
28
+
29
+ ```python
30
+ from autogen_core.memory import MemoryContent, MemoryMimeType
31
+ from goodmem_autogen import GoodMemContextProvider, GoodMemMemoryConfig
32
+
33
+ provider = GoodMemContextProvider(
34
+ config=GoodMemMemoryConfig(
35
+ base_url="https://goodmem.example.com",
36
+ api_key="gm_...", # stored as SecretStr, never serialized
37
+ space_name="handbook", # or space_id="..." to skip the lookup
38
+ embedder_id="<embedder-uuid>",
39
+ )
40
+ )
41
+
42
+ await provider.add(MemoryContent(
43
+ content="Refunds above $500 need a manager's approval.",
44
+ mime_type=MemoryMimeType.TEXT,
45
+ metadata={"title": "handbook", "category": "policy"},
46
+ ))
47
+
48
+ results = await provider.query("who approves a large refund?")
49
+ await provider.close()
50
+ ```
51
+
52
+ `add()` waits for the memory to finish indexing by default, so a query right
53
+ after it finds the result. Searching is never used as a way to wait.
54
+
55
+ Attach it to an agent and `update_context` injects retrieved passages as a
56
+ system message each turn — the same pattern AutoGen's own `ListMemory` uses.
57
+
58
+ ### What a result carries
59
+
60
+ `MemoryContent` has no score field, so provenance lives in `metadata`:
61
+
62
+ ```python
63
+ {
64
+ "title": "handbook", "category": "policy", # the memory's own metadata
65
+ "chunk_id": "...", "memory_id": "...", "space_id": "...", "source": "...",
66
+ "score": -0.53, "score_kind": "vector",
67
+ "partial": False, "statuses": [],
68
+ }
69
+ ```
70
+
71
+ - `partial` is `True` when part of the search did not complete — a reranker
72
+ was unavailable, one space was unreachable. The passages are usable but may
73
+ be incomplete, and `statuses` says why. A search that produced nothing
74
+ usable returns an empty result and emits a warning carrying the statuses,
75
+ so it is distinguishable from "no matches" without being raised. The
76
+ search tool returns `partial: true` with `statuses` in its JSON for the
77
+ same case.
78
+ - `score` is passed through exactly as GoodMem reports it. Vector scores are
79
+ opaque similarities that may be negative; reranker scores are on a scale
80
+ that depends on the reranker model (Voyage `rerank-2.5` ~`0.27..0.93`, Jina
81
+ `jina-reranker-v3` ~`-0.14..0.43` on the same documents). `score_kind` says
82
+ which you have — which is why `relevance_threshold` requires a
83
+ `reranker_id`, and why it must be calibrated for the reranker in use rather
84
+ than assumed to be 0–1. The threshold is applied by the server; if it
85
+ removes every result the provider warns, since an empty result would
86
+ otherwise read as "no matches".
87
+ - `score_kind` is read from the response, not the configuration. If the
88
+ reranker fails (`RERANKING_FAILED`, or `NOT_FOUND` for the reranker) the
89
+ server still returns the vector-search hits; they are kept, labelled
90
+ `vector`, and marked `partial` with those statuses.
91
+
92
+ ## As tools
93
+
94
+ The tools take a `goodmem.AsyncGoodmem` client, which you create and close.
95
+
96
+ ```python
97
+ from autogen_core import CancellationToken
98
+ from goodmem_autogen import create_goodmem_admin_tools, create_goodmem_search_tool
99
+ from goodmem import AsyncGoodmem
100
+
101
+ client = AsyncGoodmem(
102
+ base_url="https://goodmem.example.com",
103
+ api_key="gm_...",
104
+ # verify="/path/to/ca.pem", # a server with a self-signed certificate
105
+ )
106
+
107
+ search = create_goodmem_search_tool(
108
+ client, space_ids=["<space-uuid>"], limit=5,
109
+ reranker_id="<reranker-uuid>", # optional
110
+ metadata_filter={"category": "policy"}, # optional, escaped for you
111
+ )
112
+ admin_tools = create_goodmem_admin_tools(client) # only for agents that need them
113
+
114
+ # What an agent's tool call does:
115
+ print(await search.run_json({"query": "who approves a large refund?"}, CancellationToken()))
116
+
117
+ await client.close()
118
+ ```
119
+
120
+ The model supplies only the query; spaces, reranking and filters are yours, so
121
+ an agent cannot redirect a search or widen it mid-run. The search tool returns
122
+ JSON: `results` (each with `chunk_text`, `score`, `score_kind`, `metadata` and
123
+ IDs), `total_results`, `partial`, and `statuses` when the server reported any.
124
+
125
+ `create_goodmem_admin_tools(client)` adds space and memory management
126
+ (`goodmem_list_embedders`, `goodmem_list_rerankers`, `goodmem_list_spaces`,
127
+ `goodmem_get_space`, `goodmem_create_space`, `goodmem_update_space`,
128
+ `goodmem_delete_space`, `goodmem_create_memory`, `goodmem_list_memories`,
129
+ `goodmem_get_memory`, `goodmem_delete_memory`, and `goodmem_upload_file`
130
+ when `upload_dir` is given). A listing returns at most `max_items` (default
131
+ 100) and sets `truncated` when that cap may have cut it short. These
132
+ carry the authority of the configured API key — give them only to agents that
133
+ need them. File upload is only created when you pass `upload_dir`, and paths
134
+ resolving outside that directory are refused before the file is opened.
135
+
136
+ Every ID — a tool argument, `space_ids`, `reranker_id`, or an ID field of the
137
+ config — must be a UUID (it is lowercased); anything else raises `ValueError`
138
+ before a request is made. The SDK puts IDs into request paths unescaped, so
139
+ `"../spaces/<id>"` given as a memory ID would otherwise address a whole space.
140
+
141
+ ## Cancellation
142
+
143
+ `add`, `add_file` and `query` honour an `autogen_core.CancellationToken`: an
144
+ already-cancelled token prevents the request, and cancelling mid-flight aborts
145
+ it.
146
+
147
+ ## Clearing a space
148
+
149
+ `clear()` deletes every memory in the space and requires
150
+ `allow_clear=True` on the config, so a reflexive `clear()` cannot empty a
151
+ space by accident.
152
+
153
+ ## Filters
154
+
155
+ A filter is a GoodMem expression applied to every configured space, e.g.
156
+ `CAST(val('$.category') AS TEXT) = 'policy'`. Pass `metadata_filter={...}` and
157
+ it is built and escaped for you. Writing one by hand: inside a quoted value
158
+ escape `'` as `\'` and `\` as `\\` — SQL-style `''` doubling is rejected by
159
+ the server.
160
+
161
+ ## Development
162
+
163
+ ```bash
164
+ uv venv && uv pip install -e ".[dev]"
165
+ uv run ruff check goodmem_autogen tests
166
+ uv run mypy goodmem_autogen
167
+ uv run pytest -m "not integration" # offline: the real SDK over a mock transport or a local server
168
+ GOODMEM_BASE_URL=... GOODMEM_API_KEY=... GOODMEM_EMBEDDER_ID=... \
169
+ GOODMEM_VERIFY_SSL=false uv run pytest -m integration
170
+ ```
171
+
172
+ CI runs the first four on Python 3.10–3.13, then `uv build` and an import
173
+ of the wheel in a clean environment. It does not run the live suite, which
174
+ needs a server: `GOODMEM_EMBEDDER_ID` is the embedder its spaces are
175
+ created with, and `GOODMEM_VERIFY_SSL=false` is for a local server with a
176
+ self-signed certificate. The offline suite also executes every Python
177
+ snippet in this README against a local server.
178
+
179
+ Offline tests use event shapes captured from a live server. There is no
180
+ default API key — live tests skip unless the environment provides one.
181
+
182
+ MIT.
@@ -0,0 +1,30 @@
1
+ """goodmem-autogen: GoodMem memory and tools for the AutoGen agent framework."""
2
+
3
+ from ._config import ChunkingConfig, GoodMemMemoryConfig, PostProcessorConfig
4
+ from ._connection import GoodMemConnection
5
+ from ._context_provider import GoodMemContextProvider, GoodMemIngestionError
6
+ from ._tools import (
7
+ ADMIN_TOOL_NAMES,
8
+ SEARCH_TOOL_NAME,
9
+ create_goodmem_admin_tools,
10
+ create_goodmem_search_tool,
11
+ )
12
+ from ._uploads import GoodMemUploadError
13
+
14
+
15
+ __version__ = "0.3.0"
16
+
17
+ __all__ = [
18
+ "ADMIN_TOOL_NAMES",
19
+ "SEARCH_TOOL_NAME",
20
+ "ChunkingConfig",
21
+ "GoodMemConnection",
22
+ "GoodMemContextProvider",
23
+ "GoodMemIngestionError",
24
+ "GoodMemMemoryConfig",
25
+ "GoodMemUploadError",
26
+ "PostProcessorConfig",
27
+ "__version__",
28
+ "create_goodmem_admin_tools",
29
+ "create_goodmem_search_tool",
30
+ ]
@@ -0,0 +1,114 @@
1
+ """Configuration for the goodmem-autogen integration."""
2
+
3
+ from typing import Literal
4
+
5
+ from pydantic import BaseModel, Field, SecretStr
6
+
7
+ from ._ids import UuidStr
8
+
9
+
10
+ class PostProcessorConfig(BaseModel):
11
+ """Optional reranker / LLM post-processing applied to retrieval."""
12
+
13
+ reranker_id: UuidStr | None = Field(default=None)
14
+ llm_id: UuidStr | None = Field(default=None)
15
+ llm_temperature: float | None = Field(default=None, description="0.0-2.0")
16
+ relevance_threshold: float | None = Field(
17
+ default=None,
18
+ description=(
19
+ "Minimum reranker score, applied by the server. Only meaningful "
20
+ "with reranker_id: a raw vector score is an opaque similarity that "
21
+ "may be negative. The scale depends on the reranker model (Voyage "
22
+ "rerank-2.5 ~0.27..0.93, Jina jina-reranker-v3 ~-0.14..0.43 on the "
23
+ "same documents) and is not necessarily 0-1; calibrate it for the "
24
+ "reranker in use."
25
+ ),
26
+ )
27
+ chronological_resort: bool = Field(default=False)
28
+
29
+
30
+ class GoodMemMemoryConfig(BaseModel):
31
+ """Configuration for :class:`GoodMemContextProvider`.
32
+
33
+ ``api_key`` is a ``SecretStr``. AutoGen serializes a memory's config
34
+ whenever the owning agent is dumped (``AssistantAgent._to_config`` calls
35
+ ``memory.dump_component()``), and a plain ``str`` is written out verbatim.
36
+
37
+ ID fields must be UUIDs and are lowercased; anything else fails
38
+ validation, because the SDK puts IDs into request paths unescaped.
39
+ """
40
+
41
+ base_url: str = Field(description="GoodMem API base URL, e.g. https://goodmem.example.com")
42
+ api_key: SecretStr = Field(description="GoodMem API key (sent as X-API-Key)")
43
+ space_id: UuidStr | None = Field(
44
+ default=None,
45
+ description="Space to use (its UUID). Takes precedence over space_name.",
46
+ )
47
+ space_name: str | None = Field(
48
+ default=None,
49
+ description=(
50
+ "Space to attach to by name, created on first use if absent. "
51
+ "Reused only when its embedder matches embedder_id; an ambiguous "
52
+ "name is an error."
53
+ ),
54
+ )
55
+ embedder_id: UuidStr | None = Field(
56
+ default=None,
57
+ description="Embedder for the space (its UUID). Required when creating by space_name.",
58
+ )
59
+ max_results: int = Field(default=5, gt=0)
60
+ fetch_k: int | None = Field(
61
+ default=None, gt=0, description="Candidates to fetch before reranking."
62
+ )
63
+ filter: str | None = Field(
64
+ default=None,
65
+ description="A GoodMem filter expression applied to the space on every query.",
66
+ )
67
+ metadata: dict[str, str] | None = Field(
68
+ default=None, description="Metadata attached to everything this provider writes."
69
+ )
70
+ post_processor: PostProcessorConfig | None = Field(default=None)
71
+ upload_dir: str | None = Field(
72
+ default=None,
73
+ description=(
74
+ "Directory whose files add_file() may read. Required for file "
75
+ "uploads; paths resolving outside it are refused."
76
+ ),
77
+ )
78
+ allow_clear: bool = Field(
79
+ default=False,
80
+ description=(
81
+ "Permit clear() to delete every memory in the space. Off by "
82
+ "default so a reflexive clear() cannot empty a space."
83
+ ),
84
+ )
85
+ wait_for_indexing: bool = Field(
86
+ default=True,
87
+ description=(
88
+ "Wait for each written memory to finish indexing before add() "
89
+ "returns, so a following query finds it. Searching is never used "
90
+ "as a way to wait."
91
+ ),
92
+ )
93
+ indexing_timeout: float = Field(default=120.0, gt=0)
94
+ verify_ssl: bool = Field(default=True)
95
+ timeout: float = Field(default=60.0, gt=0)
96
+
97
+
98
+ class ChunkingConfig(BaseModel):
99
+ """Recursive chunking used when this provider creates a space."""
100
+
101
+ chunk_size: int = Field(default=512, gt=0)
102
+ chunk_overlap: int = Field(default=64, ge=0)
103
+ keep_strategy: Literal["KEEP_END", "KEEP_START", "DISCARD"] = "KEEP_END"
104
+ length_measurement: Literal["CHARACTER_COUNT", "TOKEN_COUNT"] = "CHARACTER_COUNT"
105
+
106
+ def to_api(self) -> dict[str, object]:
107
+ return {
108
+ "recursive": {
109
+ "chunkSize": self.chunk_size,
110
+ "chunkOverlap": self.chunk_overlap,
111
+ "keepStrategy": self.keep_strategy,
112
+ "lengthMeasurement": self.length_measurement,
113
+ }
114
+ }
@@ -0,0 +1,92 @@
1
+ """SDK client ownership and cancellation-aware awaiting."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ from collections.abc import Awaitable
7
+ from typing import TypeVar, cast
8
+
9
+ from autogen_core import CancellationToken
10
+ from goodmem import AsyncGoodmem
11
+ from typing_extensions import Self
12
+
13
+ from ._typing import AsyncGoodmemClient
14
+
15
+
16
+ T = TypeVar("T")
17
+
18
+
19
+ async def run_cancellable(
20
+ awaitable: Awaitable[T],
21
+ cancellation_token: CancellationToken | None = None,
22
+ ) -> T:
23
+ """Await ``awaitable``, honouring an AutoGen cancellation token.
24
+
25
+ ``CancellationToken.link_future`` cancels the wrapped future when the
26
+ token fires; this is the pattern AutoGen's own model clients use. Note
27
+ the token exposes ``is_cancelled()`` and ``link_future()`` — there is no
28
+ ``.cancelled`` attribute, despite what some ecosystem code assumes.
29
+ """
30
+ future = asyncio.ensure_future(awaitable)
31
+ if cancellation_token is not None:
32
+ if cancellation_token.is_cancelled():
33
+ future.cancel()
34
+ raise asyncio.CancelledError("operation cancelled before it was issued")
35
+ cancellation_token.link_future(future)
36
+ return await future
37
+
38
+
39
+ class GoodMemConnection:
40
+ """Owns an ``AsyncGoodmem`` client, or borrows a caller-supplied one.
41
+
42
+ An injected client keeps its own server, credentials and TLS settings;
43
+ this class never closes it. Otherwise one client is created lazily and
44
+ closed by :meth:`close`. There is no process-wide client cache.
45
+ """
46
+
47
+ def __init__(
48
+ self,
49
+ *,
50
+ base_url: str | None = None,
51
+ api_key: str | None = None,
52
+ verify_ssl: bool | str = True,
53
+ timeout: float = 60.0,
54
+ client: AsyncGoodmem | None = None,
55
+ ) -> None:
56
+ self._base_url = base_url
57
+ self._api_key = api_key
58
+ self._verify_ssl = verify_ssl
59
+ self._timeout = timeout
60
+ self._injected = client
61
+ self._owned: AsyncGoodmem | None = None
62
+
63
+ @property
64
+ def owns_client(self) -> bool:
65
+ return self._injected is None
66
+
67
+ def client(self) -> AsyncGoodmemClient:
68
+ if self._injected is not None:
69
+ return cast(AsyncGoodmemClient, self._injected)
70
+ if self._owned is None:
71
+ if not self._base_url or not self._api_key:
72
+ raise ValueError(
73
+ "GoodMem base_url and api_key are required when no client is injected."
74
+ )
75
+ self._owned = AsyncGoodmem(
76
+ base_url=self._base_url.rstrip("/"),
77
+ api_key=self._api_key,
78
+ verify=self._verify_ssl,
79
+ timeout=self._timeout,
80
+ )
81
+ return cast(AsyncGoodmemClient, self._owned)
82
+
83
+ async def close(self) -> None:
84
+ if self._owned is not None:
85
+ await self._owned.close()
86
+ self._owned = None
87
+
88
+ async def __aenter__(self) -> Self:
89
+ return self
90
+
91
+ async def __aexit__(self, *exc: object) -> None:
92
+ await self.close()