agently 4.0.6.5__tar.gz → 4.0.6.7__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 (77) hide show
  1. {agently-4.0.6.5 → agently-4.0.6.7}/PKG-INFO +1 -1
  2. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/plugins/ModelRequester/OpenAICompatible.py +3 -3
  3. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/plugins/PromptGenerator/AgentlyPromptGenerator.py +4 -16
  4. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/ModelRequest.py +1 -2
  5. agently-4.0.6.7/agently/integrations/chromadb.py +283 -0
  6. {agently-4.0.6.5 → agently-4.0.6.7}/pyproject.toml +1 -1
  7. agently-4.0.6.5/agently/integrations/chromadb.py +0 -167
  8. {agently-4.0.6.5 → agently-4.0.6.7}/LICENSE +0 -0
  9. {agently-4.0.6.5 → agently-4.0.6.7}/README.md +0 -0
  10. {agently-4.0.6.5 → agently-4.0.6.7}/agently/__init__.py +0 -0
  11. {agently-4.0.6.5 → agently-4.0.6.7}/agently/_default_init.py +0 -0
  12. {agently-4.0.6.5 → agently-4.0.6.7}/agently/_default_settings.yaml +0 -0
  13. {agently-4.0.6.5 → agently-4.0.6.7}/agently/base.py +0 -0
  14. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/agent_extensions/AutoFuncExtension.py +0 -0
  15. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/agent_extensions/ChatSessionExtension.py +0 -0
  16. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/agent_extensions/ConfigurePromptExtension.py +0 -0
  17. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/agent_extensions/KeyWaiterExtension.py +0 -0
  18. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/agent_extensions/ToolExtension.py +0 -0
  19. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/agent_extensions/__init__.py +0 -0
  20. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/hookers/ConsoleHooker.py +0 -0
  21. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/hookers/PureLoggerHooker.py +0 -0
  22. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/hookers/SystemMessageHooker.py +0 -0
  23. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/plugins/ResponseParser/AgentlyResponseParser.py +0 -0
  24. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/plugins/ToolManager/AgentlyToolManager.py +0 -0
  25. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/plugins/__init__.py +0 -0
  26. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/tools/Browse.py +0 -0
  27. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/tools/Search.py +0 -0
  28. {agently-4.0.6.5 → agently-4.0.6.7}/agently/builtins/tools/__init__.py +0 -0
  29. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/Agent.py +0 -0
  30. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/EventCenter.py +0 -0
  31. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/ExtensionHandlers.py +0 -0
  32. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/PluginManager.py +0 -0
  33. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/Prompt.py +0 -0
  34. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/Tool.py +0 -0
  35. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/BluePrint.py +0 -0
  36. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/Chunk.py +0 -0
  37. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/Execution.py +0 -0
  38. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/Process.py +0 -0
  39. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/TriggerFlow.py +0 -0
  40. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/__init__.py +0 -0
  41. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/process/BaseProcess.py +0 -0
  42. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/process/ForEachProcess.py +0 -0
  43. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/process/MatchCaseProcess.py +0 -0
  44. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/TriggerFlow/process/__init__.py +0 -0
  45. {agently-4.0.6.5 → agently-4.0.6.7}/agently/core/__init__.py +0 -0
  46. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/__init__.py +0 -0
  47. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/data/__init__.py +0 -0
  48. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/data/event.py +0 -0
  49. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/data/prompt.py +0 -0
  50. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/data/request.py +0 -0
  51. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/data/response.py +0 -0
  52. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/data/serializable.py +0 -0
  53. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/data/tool.py +0 -0
  54. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/plugins/EventHooker.py +0 -0
  55. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/plugins/ModelRequester.py +0 -0
  56. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/plugins/PromptGenerator.py +0 -0
  57. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/plugins/ResponseParser.py +0 -0
  58. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/plugins/ToolManager.py +0 -0
  59. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/plugins/__init__.py +0 -0
  60. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/plugins/base.py +0 -0
  61. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/trigger_flow/__init__.py +0 -0
  62. {agently-4.0.6.5 → agently-4.0.6.7}/agently/types/trigger_flow/trigger_flow.py +0 -0
  63. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/DataFormatter.py +0 -0
  64. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/DataLocator.py +0 -0
  65. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/DataPathBuilder.py +0 -0
  66. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/FunctionShifter.py +0 -0
  67. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/GeneratorConsumer.py +0 -0
  68. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/LazyImport.py +0 -0
  69. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/Logger.py +0 -0
  70. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/Messenger.py +0 -0
  71. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/RuntimeData.py +0 -0
  72. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/SerializableRuntimeData.py +0 -0
  73. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/Settings.py +0 -0
  74. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/Storage.py +0 -0
  75. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/StreamingJSONCompleter.py +0 -0
  76. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/StreamingJSONParser.py +0 -0
  77. {agently-4.0.6.5 → agently-4.0.6.7}/agently/utils/__init__.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: agently
3
- Version: 4.0.6.5
3
+ Version: 4.0.6.7
4
4
  Summary:
5
5
  License: Apache-2.0
6
6
  License-File: LICENSE
@@ -351,13 +351,13 @@ class OpenAICompatible(ModelRequester):
351
351
  api_key = self.plugin_settings.get("api_key", None)
352
352
  if api_key is not None and auth["api_key"] == "None":
353
353
  auth["api_key"] = str(api_key)
354
- if "api_key" in auth:
355
- headers_with_auth = {**request_data.headers, "Authorization": f"Bearer { auth['api_key'] }"}
356
- elif "headers" in auth and isinstance(auth["headers"], dict):
354
+ if "headers" in auth and isinstance(auth["headers"], dict):
357
355
  headers_with_auth = {**request_data.headers, **auth["headers"]}
358
356
  elif "body" in auth and isinstance(auth["body"], dict):
359
357
  headers_with_auth = request_data.headers.copy()
360
358
  request_data.data.update(**auth["body"])
359
+ if "api_key" in auth and auth["api_key"] != "None":
360
+ headers_with_auth = {**request_data.headers, "Authorization": f"Bearer { auth['api_key'] }"}
361
361
  else:
362
362
  headers_with_auth = request_data.headers.copy()
363
363
 
@@ -307,9 +307,7 @@ class AgentlyPromptGenerator(PromptGenerator):
307
307
  if isinstance(role_mapping, dict):
308
308
  merged_role_mapping.update(role_mapping)
309
309
 
310
- prompt_text_list.append(
311
- f"{ (merged_role_mapping['user'] if 'user' in merged_role_mapping else 'user').upper() }:"
312
- )
310
+ prompt_text_list.append(f"{ (merged_role_mapping['user'] if 'user' in merged_role_mapping else 'user') }:")
313
311
 
314
312
  # system & developer
315
313
  if prompt_object.system:
@@ -368,7 +366,7 @@ class AgentlyPromptGenerator(PromptGenerator):
368
366
 
369
367
  prompt_text_list.extend(self._generate_main_prompt(prompt_object))
370
368
  prompt_text_list.append(
371
- f"{ (merged_role_mapping['assistant'] if 'assistant' in merged_role_mapping else 'assistant').upper() }:"
369
+ f"{ (merged_role_mapping['assistant'] if 'assistant' in merged_role_mapping else 'assistant') }:"
372
370
  )
373
371
 
374
372
  return "\n".join(prompt_text_list)
@@ -398,12 +396,7 @@ class AgentlyPromptGenerator(PromptGenerator):
398
396
  if prompt_object.system:
399
397
  prompt_messages.append(
400
398
  self._generate_yaml_prompt_message(
401
- str(
402
- prompt_title_mapping.get(
403
- 'system',
404
- 'SYSTEM',
405
- )
406
- ),
399
+ "system",
407
400
  prompt_object.system,
408
401
  role_mapping=merged_role_mapping,
409
402
  )
@@ -412,12 +405,7 @@ class AgentlyPromptGenerator(PromptGenerator):
412
405
  if prompt_object.developer:
413
406
  prompt_messages.append(
414
407
  self._generate_yaml_prompt_message(
415
- str(
416
- prompt_title_mapping.get(
417
- 'developer',
418
- 'DEVELOPER DIRECTIONS',
419
- )
420
- ),
408
+ "developer",
421
409
  prompt_object.developer,
422
410
  role_mapping=merged_role_mapping,
423
411
  )
@@ -99,7 +99,6 @@ class ModelResponse:
99
99
  self.plugin_manager,
100
100
  self.settings,
101
101
  )
102
- self.get_result = self.result
103
102
  self.get_meta = self.result.get_meta
104
103
  self.async_get_meta = self.result.async_get_meta
105
104
  self.get_text = self.result.get_text
@@ -406,7 +405,7 @@ class ModelRequest:
406
405
  self.prompt.clear()
407
406
  return response
408
407
 
409
- async def get_result(self):
408
+ def get_result(self):
410
409
  return self.get_response().result
411
410
 
412
411
  async def async_get_meta(self):
@@ -0,0 +1,283 @@
1
+ from agently.utils import LazyImport
2
+
3
+ LazyImport.import_package("chromadb")
4
+
5
+ from typing import Literal, Callable, TypedDict, Any, TYPE_CHECKING, overload
6
+ from itertools import zip_longest
7
+
8
+ from chromadb import Client as ChromaDBClient
9
+ from chromadb.config import Settings
10
+ from chromadb.api.types import EmbeddingFunction
11
+
12
+ if TYPE_CHECKING:
13
+ from chromadb.api.types import (
14
+ Documents,
15
+ Embedding,
16
+ Embeddings,
17
+ QueryResult,
18
+ Schema,
19
+ DataLoader,
20
+ Loadable,
21
+ Where,
22
+ WhereDocument,
23
+ )
24
+ from chromadb.api.collection_configuration import CreateCollectionConfiguration
25
+ from chromadb.api import ClientAPI
26
+ from agently.core import BaseAgent
27
+
28
+
29
+ class ChromaDataDictOptional(TypedDict, total=False):
30
+ metadata: dict[Any, Any]
31
+ id: Any
32
+ embedding: "Embedding"
33
+
34
+
35
+ class ChromaDataDict(ChromaDataDictOptional):
36
+ document: str
37
+
38
+
39
+ class ChromaData:
40
+ def __init__(
41
+ self,
42
+ original_data: ChromaDataDict | list[ChromaDataDict],
43
+ *,
44
+ embedding_function: "Callable[[str | list[str]], Embeddings] | None" = None,
45
+ agent: "BaseAgent | None" = None,
46
+ ):
47
+ self._original_data = original_data if isinstance(original_data, list) else [original_data]
48
+ if embedding_function:
49
+ self._embedding_function = embedding_function
50
+ elif agent:
51
+
52
+ def embedding_function_by_agent(texts: str | list[str]) -> "Embeddings":
53
+ return agent.input(texts).start()
54
+
55
+ self._embedding_function = embedding_function_by_agent
56
+ else:
57
+ self._embedding_function = None
58
+ self._chroma_data = {
59
+ "documents": [],
60
+ "metadatas": [],
61
+ "ids": [],
62
+ "embeddings": [],
63
+ }
64
+ if self._embedding_function is not None:
65
+ documents = [data["document"] for data in self._original_data]
66
+ embeddings = self._embedding_function(documents)
67
+ self._chroma_data["embeddings"] = embeddings
68
+ for index, data in enumerate(self._original_data):
69
+ data["embedding"] = embeddings[index]
70
+ else:
71
+ del self._chroma_data["embeddings"]
72
+ for index, data in enumerate(self._original_data):
73
+ self._chroma_data["documents"].append(data["document"])
74
+ self._chroma_data["metadatas"].append(data["metadata"] if "metadata" in data else None)
75
+ self._chroma_data["ids"].append(data["id"] if "id" in data else str(index))
76
+
77
+ def get_kwargs(self):
78
+ return self._chroma_data
79
+
80
+ def get_original_data(self):
81
+ return self._original_data
82
+
83
+
84
+ class ChromaResults:
85
+ def __init__(
86
+ self,
87
+ *,
88
+ queries: str | list[str],
89
+ results: "QueryResult",
90
+ distance: float | None = None,
91
+ ):
92
+ if isinstance(queries, str):
93
+ queries = [queries]
94
+
95
+ ids = results.get("ids") or []
96
+ documents = results.get("documents")
97
+ metadatas = results.get("metadatas")
98
+ distances = results.get("distances")
99
+
100
+ if ids and len(queries) != len(ids):
101
+ raise ValueError(
102
+ f"The length of queries does not equal the length of results['ids'].\nQueries: {queries}\nIds: {ids}"
103
+ )
104
+
105
+ self._chroma_results = results
106
+ self._results: dict[str, list[dict]] = {}
107
+
108
+ doc_iter = documents if documents is not None else []
109
+ meta_iter = metadatas if metadatas is not None else []
110
+ dist_iter = distances if distances is not None else []
111
+
112
+ for query, row_ids, row_documents, row_metadatas, row_distances in zip_longest(
113
+ queries, ids, doc_iter, meta_iter, dist_iter, fillvalue=None
114
+ ):
115
+ row_ids = row_ids or []
116
+ query_results: list[dict[str, Any]] = []
117
+ for id_, doc, meta, dist in zip_longest(
118
+ row_ids, row_documents or [], row_metadatas or [], row_distances or [], fillvalue=None
119
+ ):
120
+ query_result = {"id": id_, "document": doc, "metadata": meta, "distance": dist}
121
+ if distance is None:
122
+ query_results.append(query_result)
123
+ else:
124
+ if isinstance(dist, (int, float)) and dist < distance:
125
+ query_results.append(query_result)
126
+
127
+ if query:
128
+ self._results[query] = query_results
129
+
130
+ def get_original_results(self):
131
+ return self._chroma_results
132
+
133
+ def get(self, *, simplify_single_result: bool = False):
134
+ if simplify_single_result:
135
+ results = list(self._results.values())
136
+ if len(results) == 1:
137
+ return results[0]
138
+ return self._results
139
+
140
+
141
+ class ChromaEmbeddingFunction(EmbeddingFunction):
142
+ def __init__(
143
+ self,
144
+ *,
145
+ embedding_agent: "BaseAgent",
146
+ ):
147
+ def embedding_function_by_agent(texts: list[str]) -> "Embeddings":
148
+ return embedding_agent.input(texts).start()
149
+
150
+ self.embedding_function = embedding_function_by_agent
151
+
152
+ def __call__(self, documents: "Documents | list[str] | str") -> "Embeddings":
153
+ if isinstance(documents, str):
154
+ documents = [documents]
155
+ return self.embedding_function([document for document in documents])
156
+
157
+
158
+ class ChromaCollection:
159
+
160
+ def __init__(
161
+ self,
162
+ collection_name: str,
163
+ *,
164
+ conn: "ClientAPI | None" = None,
165
+ schema: "Schema | None" = None,
166
+ configuration: "CreateCollectionConfiguration | None" = None,
167
+ metadata: dict[str, Any] | None = None,
168
+ embedding_agent: "BaseAgent | None" = None,
169
+ data_loader: "DataLoader[Loadable] | None" = None,
170
+ get_or_create: bool = False,
171
+ hnsw_space: Literal["l2", "cosine", "ip"] = "cosine",
172
+ ):
173
+ if conn is None:
174
+ conn = ChromaDBClient(Settings(anonymized_telemetry=False))
175
+
176
+ if metadata is None:
177
+ metadata = {}
178
+
179
+ metadata["hnsw:space"] = hnsw_space
180
+
181
+ self.embedding_function = (
182
+ ChromaEmbeddingFunction(embedding_agent=embedding_agent) if embedding_agent is not None else None
183
+ )
184
+
185
+ self._collection = conn.create_collection(
186
+ collection_name,
187
+ schema=schema,
188
+ configuration=configuration,
189
+ metadata=metadata,
190
+ embedding_function=self.embedding_function,
191
+ data_loader=data_loader,
192
+ get_or_create=get_or_create,
193
+ )
194
+
195
+ def add(
196
+ self,
197
+ data: ChromaDataDict | list[ChromaDataDict] | ChromaData,
198
+ *,
199
+ embedding_function: "Callable[[str | list[str]], Embeddings] | None" = None,
200
+ ):
201
+ embedding_function = embedding_function or self.embedding_function
202
+ if embedding_function is None:
203
+ raise NotImplementedError(
204
+ f"Embedding function must be assigned either in ChromaCollection initialize or .add() method."
205
+ )
206
+ if not isinstance(data, ChromaData):
207
+ data = ChromaData(data, embedding_function=embedding_function)
208
+ self._collection.add(**data.get_kwargs())
209
+ return self
210
+
211
+ @overload
212
+ def query(
213
+ self,
214
+ query_or_queries: str,
215
+ *,
216
+ embedding_function: "Callable[[str | list[str]], Embeddings] | None" = None,
217
+ top_n: int | None = None,
218
+ where: "Where | None" = None,
219
+ where_document: "WhereDocument | None" = None,
220
+ distance: float | None = None,
221
+ ) -> list[dict[str, Any]]: ...
222
+
223
+ @overload
224
+ def query(
225
+ self,
226
+ query_or_queries: list[str],
227
+ *,
228
+ embedding_function: "Callable[[str | list[str]], Embeddings] | None" = None,
229
+ top_n: int | None = None,
230
+ where: "Where | None" = None,
231
+ where_document: "WhereDocument | None" = None,
232
+ distance: float | None = None,
233
+ ) -> dict[str, list[dict[str, Any]]]: ...
234
+
235
+ def query(
236
+ self,
237
+ query_or_queries: str | list[str],
238
+ *,
239
+ embedding_function: "Callable[[str | list[str]], Embeddings] | None" = None,
240
+ top_n: int | None = None,
241
+ where: "Where | None" = None,
242
+ where_document: "WhereDocument | None" = None,
243
+ distance: float | None = None,
244
+ ):
245
+ embedding_function = embedding_function or self.embedding_function
246
+ if embedding_function is None:
247
+ raise NotImplementedError(
248
+ f"Embedding function must be assigned either in ChromaCollection initialize or .query() method."
249
+ )
250
+ if not isinstance(query_or_queries, list):
251
+ query_or_queries = [query_or_queries]
252
+ return ChromaResults(
253
+ queries=query_or_queries,
254
+ results=self._collection.query(
255
+ query_embeddings=embedding_function(query_or_queries),
256
+ n_results=top_n or 10,
257
+ where=where,
258
+ where_document=where_document,
259
+ ),
260
+ distance=distance,
261
+ ).get(simplify_single_result=True)
262
+
263
+ def query_embeddings(
264
+ self,
265
+ embeddings_dict: dict[str, "Embedding"],
266
+ *,
267
+ top_n: int | None = None,
268
+ where: "Where | None" = None,
269
+ where_document: "WhereDocument | None" = None,
270
+ distance: float | None = None,
271
+ ):
272
+ queries = [query for query in embeddings_dict.keys()]
273
+ embeddings = [embeddings for embeddings in embeddings_dict.values()]
274
+ return ChromaResults(
275
+ queries=queries,
276
+ results=self._collection.query(
277
+ query_embeddings=embeddings,
278
+ n_results=top_n or 10,
279
+ where=where,
280
+ where_document=where_document,
281
+ ),
282
+ distance=distance,
283
+ ).get(simplify_single_result=True)
@@ -1,6 +1,6 @@
1
1
  [project]
2
2
  name = "agently"
3
- version = "4.0.6.5"
3
+ version = "4.0.6.7"
4
4
  description = ""
5
5
  authors = [
6
6
  {name = "Agently Team",email = "developer@agently.tech"},
@@ -1,167 +0,0 @@
1
- from agently.utils import LazyImport, FunctionShifter
2
-
3
- LazyImport.import_package("chromadb")
4
-
5
- from typing import Callable, TypedDict, Any, TYPE_CHECKING
6
-
7
- from chromadb.api.types import EmbeddingFunction, DefaultEmbeddingFunction
8
-
9
- if TYPE_CHECKING:
10
- from chromadb.api.types import (
11
- Documents,
12
- Embedding,
13
- Embeddings,
14
- QueryResult,
15
- Schema,
16
- CollectionMetadata,
17
- DataLoader,
18
- Loadable,
19
- )
20
- from chromadb.api.collection_configuration import CreateCollectionConfiguration
21
- from chromadb.api import ClientAPI
22
- from agently.core import BaseAgent
23
-
24
-
25
- class ChromaDataDictOptional(TypedDict, total=False):
26
- metadata: dict[Any, Any]
27
- id: Any
28
- embedding: "Embedding"
29
-
30
-
31
- class ChromaDataDict(ChromaDataDictOptional):
32
- document: str
33
-
34
-
35
- class ChromaData:
36
- def __init__(
37
- self,
38
- original_data: ChromaDataDict | list[ChromaDataDict],
39
- *,
40
- embedding_function: "Callable[[str | list[str]], Embeddings] | None" = None,
41
- agent: "BaseAgent | None" = None,
42
- ):
43
- self._original_data = original_data if isinstance(original_data, list) else [original_data]
44
- if embedding_function:
45
- self._embedding_function = embedding_function
46
- elif agent:
47
-
48
- def embedding_function_by_agent(texts: str | list[str]) -> "Embeddings":
49
- return agent.input(texts).start()
50
-
51
- self._embedding_function = embedding_function_by_agent
52
- else:
53
- self._embedding_function = None
54
- self._chroma_data = {
55
- "documents": [],
56
- "metadatas": [],
57
- "ids": [],
58
- "embeddings": [],
59
- }
60
- if self._embedding_function is not None:
61
- documents = [data["document"] for data in self._original_data]
62
- embeddings = self._embedding_function(documents)
63
- self._chroma_data["embeddings"] = embeddings
64
- for index, data in enumerate(self._original_data):
65
- data["embedding"] = embeddings[index]
66
- else:
67
- del self._chroma_data["embeddings"]
68
- for index, data in enumerate(self._original_data):
69
- self._chroma_data["documents"].append(data["document"])
70
- self._chroma_data["metadatas"].append(data["metadata"] if "metadata" in data else None)
71
- self._chroma_data["ids"].append(data["id"] if "id" in data else str(index))
72
-
73
- def get_kwargs(self):
74
- return self._chroma_data
75
-
76
- def get_original_data(self):
77
- return self._original_data
78
-
79
-
80
- class ChromaResults:
81
- def __init__(
82
- self,
83
- *,
84
- queries: str | list[str],
85
- results: "QueryResult",
86
- ):
87
- if isinstance(queries, str):
88
- queries = [queries]
89
- if len(queries) != len(results["documents"] or []):
90
- raise ValueError(
91
- f"The length of queries does not equal the length of results.\nQueries: { queries }\nResults: { results }"
92
- )
93
- self._chroma_results = results
94
- self._results = {}
95
- for query_index, query in enumerate(queries):
96
- for result_index, id in enumerate(results["ids"][query_index]):
97
- result = {
98
- "id": id,
99
- "document": results["documents"][query_index][result_index] if results["documents"] else None,
100
- "metadata": results["metadatas"][query_index][result_index] if results["metadatas"] else None,
101
- "distance": results["distances"][query_index][result_index] if results["distances"] else None,
102
- }
103
- self._results[query] = result
104
-
105
- def get_original_results(self):
106
- return self._chroma_results
107
-
108
- def get(self):
109
- return self._results
110
-
111
-
112
- class ChromaEmbeddingFunction(EmbeddingFunction):
113
- def __init__(
114
- self,
115
- *,
116
- embedding_function: "Callable[[list[str]], Embeddings] | None" = None,
117
- agent: "BaseAgent | None" = None,
118
- ):
119
- if embedding_function:
120
- self.embedding_function = embedding_function
121
- elif agent:
122
-
123
- def embedding_function_by_agent(texts: list[str]) -> "Embeddings":
124
- return agent.input(texts).start()
125
-
126
- self.embedding_function = embedding_function_by_agent
127
- else:
128
- raise ValueError(
129
- f"ChromaEmbeddingFunction() requires at least one definition for 'embedding_function' or 'agent'."
130
- )
131
-
132
- def __call__(self, documents: "Documents") -> "Embeddings":
133
- return self.embedding_function([document for document in documents])
134
-
135
-
136
- class ChromaCollection:
137
- def __init__(
138
- self,
139
- conn: "ClientAPI",
140
- collection_name: str,
141
- *,
142
- schema: "Schema | None" = None,
143
- configuration: "CreateCollectionConfiguration | None" = None,
144
- metadata: None = None,
145
- embedding_function: EmbeddingFunction | None = None,
146
- agent: "BaseAgent | None" = None,
147
- data_loader: "DataLoader[Loadable] | None" = None,
148
- get_or_create: bool = False,
149
- ):
150
- if embedding_function:
151
- self._embedding_function = embedding_function
152
- elif agent:
153
- self._embedding_function = ChromaEmbeddingFunction(agent=agent)
154
- else:
155
- self._embedding_function = DefaultEmbeddingFunction()
156
-
157
- self._collection = conn.create_collection(
158
- collection_name,
159
- schema=schema,
160
- configuration=configuration,
161
- metadata=metadata,
162
- embedding_function=embedding_function,
163
- data_loader=data_loader,
164
- get_or_create=get_or_create,
165
- )
166
-
167
- def add(self): ...
File without changes
File without changes
File without changes
File without changes