agently 4.0.6.6__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.6 → agently-4.0.6.7}/PKG-INFO +1 -1
  2. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/plugins/ModelRequester/OpenAICompatible.py +3 -3
  3. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/ModelRequest.py +1 -2
  4. agently-4.0.6.7/agently/integrations/chromadb.py +283 -0
  5. {agently-4.0.6.6 → agently-4.0.6.7}/pyproject.toml +1 -1
  6. agently-4.0.6.6/agently/integrations/chromadb.py +0 -167
  7. {agently-4.0.6.6 → agently-4.0.6.7}/LICENSE +0 -0
  8. {agently-4.0.6.6 → agently-4.0.6.7}/README.md +0 -0
  9. {agently-4.0.6.6 → agently-4.0.6.7}/agently/__init__.py +0 -0
  10. {agently-4.0.6.6 → agently-4.0.6.7}/agently/_default_init.py +0 -0
  11. {agently-4.0.6.6 → agently-4.0.6.7}/agently/_default_settings.yaml +0 -0
  12. {agently-4.0.6.6 → agently-4.0.6.7}/agently/base.py +0 -0
  13. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/agent_extensions/AutoFuncExtension.py +0 -0
  14. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/agent_extensions/ChatSessionExtension.py +0 -0
  15. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/agent_extensions/ConfigurePromptExtension.py +0 -0
  16. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/agent_extensions/KeyWaiterExtension.py +0 -0
  17. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/agent_extensions/ToolExtension.py +0 -0
  18. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/agent_extensions/__init__.py +0 -0
  19. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/hookers/ConsoleHooker.py +0 -0
  20. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/hookers/PureLoggerHooker.py +0 -0
  21. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/hookers/SystemMessageHooker.py +0 -0
  22. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/plugins/PromptGenerator/AgentlyPromptGenerator.py +0 -0
  23. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/plugins/ResponseParser/AgentlyResponseParser.py +0 -0
  24. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/plugins/ToolManager/AgentlyToolManager.py +0 -0
  25. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/plugins/__init__.py +0 -0
  26. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/tools/Browse.py +0 -0
  27. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/tools/Search.py +0 -0
  28. {agently-4.0.6.6 → agently-4.0.6.7}/agently/builtins/tools/__init__.py +0 -0
  29. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/Agent.py +0 -0
  30. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/EventCenter.py +0 -0
  31. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/ExtensionHandlers.py +0 -0
  32. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/PluginManager.py +0 -0
  33. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/Prompt.py +0 -0
  34. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/Tool.py +0 -0
  35. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/BluePrint.py +0 -0
  36. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/Chunk.py +0 -0
  37. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/Execution.py +0 -0
  38. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/Process.py +0 -0
  39. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/TriggerFlow.py +0 -0
  40. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/__init__.py +0 -0
  41. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/process/BaseProcess.py +0 -0
  42. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/process/ForEachProcess.py +0 -0
  43. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/process/MatchCaseProcess.py +0 -0
  44. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/TriggerFlow/process/__init__.py +0 -0
  45. {agently-4.0.6.6 → agently-4.0.6.7}/agently/core/__init__.py +0 -0
  46. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/__init__.py +0 -0
  47. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/data/__init__.py +0 -0
  48. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/data/event.py +0 -0
  49. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/data/prompt.py +0 -0
  50. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/data/request.py +0 -0
  51. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/data/response.py +0 -0
  52. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/data/serializable.py +0 -0
  53. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/data/tool.py +0 -0
  54. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/plugins/EventHooker.py +0 -0
  55. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/plugins/ModelRequester.py +0 -0
  56. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/plugins/PromptGenerator.py +0 -0
  57. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/plugins/ResponseParser.py +0 -0
  58. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/plugins/ToolManager.py +0 -0
  59. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/plugins/__init__.py +0 -0
  60. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/plugins/base.py +0 -0
  61. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/trigger_flow/__init__.py +0 -0
  62. {agently-4.0.6.6 → agently-4.0.6.7}/agently/types/trigger_flow/trigger_flow.py +0 -0
  63. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/DataFormatter.py +0 -0
  64. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/DataLocator.py +0 -0
  65. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/DataPathBuilder.py +0 -0
  66. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/FunctionShifter.py +0 -0
  67. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/GeneratorConsumer.py +0 -0
  68. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/LazyImport.py +0 -0
  69. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/Logger.py +0 -0
  70. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/Messenger.py +0 -0
  71. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/RuntimeData.py +0 -0
  72. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/SerializableRuntimeData.py +0 -0
  73. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/Settings.py +0 -0
  74. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/Storage.py +0 -0
  75. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/StreamingJSONCompleter.py +0 -0
  76. {agently-4.0.6.6 → agently-4.0.6.7}/agently/utils/StreamingJSONParser.py +0 -0
  77. {agently-4.0.6.6 → 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.6
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
 
@@ -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.6"
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