ag2 0.4.1__py3-none-any.whl → 0.4.2b1__py3-none-any.whl

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.

Potentially problematic release.


This version of ag2 might be problematic. Click here for more details.

Files changed (161) hide show
  1. ag2-0.4.2b1.dist-info/METADATA +19 -0
  2. ag2-0.4.2b1.dist-info/RECORD +6 -0
  3. ag2-0.4.2b1.dist-info/top_level.txt +1 -0
  4. ag2-0.4.1.dist-info/METADATA +0 -500
  5. ag2-0.4.1.dist-info/RECORD +0 -158
  6. ag2-0.4.1.dist-info/top_level.txt +0 -1
  7. autogen/__init__.py +0 -17
  8. autogen/_pydantic.py +0 -116
  9. autogen/agentchat/__init__.py +0 -42
  10. autogen/agentchat/agent.py +0 -142
  11. autogen/agentchat/assistant_agent.py +0 -85
  12. autogen/agentchat/chat.py +0 -306
  13. autogen/agentchat/contrib/__init__.py +0 -0
  14. autogen/agentchat/contrib/agent_builder.py +0 -788
  15. autogen/agentchat/contrib/agent_eval/agent_eval.py +0 -107
  16. autogen/agentchat/contrib/agent_eval/criterion.py +0 -47
  17. autogen/agentchat/contrib/agent_eval/critic_agent.py +0 -47
  18. autogen/agentchat/contrib/agent_eval/quantifier_agent.py +0 -42
  19. autogen/agentchat/contrib/agent_eval/subcritic_agent.py +0 -48
  20. autogen/agentchat/contrib/agent_eval/task.py +0 -43
  21. autogen/agentchat/contrib/agent_optimizer.py +0 -450
  22. autogen/agentchat/contrib/capabilities/__init__.py +0 -0
  23. autogen/agentchat/contrib/capabilities/agent_capability.py +0 -21
  24. autogen/agentchat/contrib/capabilities/generate_images.py +0 -297
  25. autogen/agentchat/contrib/capabilities/teachability.py +0 -406
  26. autogen/agentchat/contrib/capabilities/text_compressors.py +0 -72
  27. autogen/agentchat/contrib/capabilities/transform_messages.py +0 -92
  28. autogen/agentchat/contrib/capabilities/transforms.py +0 -565
  29. autogen/agentchat/contrib/capabilities/transforms_util.py +0 -120
  30. autogen/agentchat/contrib/capabilities/vision_capability.py +0 -217
  31. autogen/agentchat/contrib/captainagent/tools/__init__.py +0 -0
  32. autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_correlation.py +0 -41
  33. autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_skewness_and_kurtosis.py +0 -29
  34. autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_iqr.py +0 -29
  35. autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_zscore.py +0 -29
  36. autogen/agentchat/contrib/captainagent/tools/data_analysis/explore_csv.py +0 -22
  37. autogen/agentchat/contrib/captainagent/tools/data_analysis/shapiro_wilk_test.py +0 -31
  38. autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_download.py +0 -26
  39. autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_search.py +0 -55
  40. autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_image.py +0 -54
  41. autogen/agentchat/contrib/captainagent/tools/information_retrieval/extract_pdf_text.py +0 -39
  42. autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_wikipedia_text.py +0 -22
  43. autogen/agentchat/contrib/captainagent/tools/information_retrieval/get_youtube_caption.py +0 -35
  44. autogen/agentchat/contrib/captainagent/tools/information_retrieval/image_qa.py +0 -61
  45. autogen/agentchat/contrib/captainagent/tools/information_retrieval/optical_character_recognition.py +0 -62
  46. autogen/agentchat/contrib/captainagent/tools/information_retrieval/perform_web_search.py +0 -48
  47. autogen/agentchat/contrib/captainagent/tools/information_retrieval/scrape_wikipedia_tables.py +0 -34
  48. autogen/agentchat/contrib/captainagent/tools/information_retrieval/transcribe_audio_file.py +0 -22
  49. autogen/agentchat/contrib/captainagent/tools/information_retrieval/youtube_download.py +0 -36
  50. autogen/agentchat/contrib/captainagent/tools/math/calculate_circle_area_from_diameter.py +0 -22
  51. autogen/agentchat/contrib/captainagent/tools/math/calculate_day_of_the_week.py +0 -19
  52. autogen/agentchat/contrib/captainagent/tools/math/calculate_fraction_sum.py +0 -29
  53. autogen/agentchat/contrib/captainagent/tools/math/calculate_matrix_power.py +0 -32
  54. autogen/agentchat/contrib/captainagent/tools/math/calculate_reflected_point.py +0 -17
  55. autogen/agentchat/contrib/captainagent/tools/math/complex_numbers_product.py +0 -26
  56. autogen/agentchat/contrib/captainagent/tools/math/compute_currency_conversion.py +0 -24
  57. autogen/agentchat/contrib/captainagent/tools/math/count_distinct_permutations.py +0 -28
  58. autogen/agentchat/contrib/captainagent/tools/math/evaluate_expression.py +0 -29
  59. autogen/agentchat/contrib/captainagent/tools/math/find_continuity_point.py +0 -35
  60. autogen/agentchat/contrib/captainagent/tools/math/fraction_to_mixed_numbers.py +0 -40
  61. autogen/agentchat/contrib/captainagent/tools/math/modular_inverse_sum.py +0 -23
  62. autogen/agentchat/contrib/captainagent/tools/math/simplify_mixed_numbers.py +0 -37
  63. autogen/agentchat/contrib/captainagent/tools/math/sum_of_digit_factorials.py +0 -16
  64. autogen/agentchat/contrib/captainagent/tools/math/sum_of_primes_below.py +0 -16
  65. autogen/agentchat/contrib/captainagent/tools/requirements.txt +0 -10
  66. autogen/agentchat/contrib/captainagent/tools/tool_description.tsv +0 -34
  67. autogen/agentchat/contrib/captainagent.py +0 -490
  68. autogen/agentchat/contrib/gpt_assistant_agent.py +0 -545
  69. autogen/agentchat/contrib/graph_rag/__init__.py +0 -0
  70. autogen/agentchat/contrib/graph_rag/document.py +0 -30
  71. autogen/agentchat/contrib/graph_rag/falkor_graph_query_engine.py +0 -111
  72. autogen/agentchat/contrib/graph_rag/falkor_graph_rag_capability.py +0 -81
  73. autogen/agentchat/contrib/graph_rag/graph_query_engine.py +0 -56
  74. autogen/agentchat/contrib/graph_rag/graph_rag_capability.py +0 -64
  75. autogen/agentchat/contrib/img_utils.py +0 -390
  76. autogen/agentchat/contrib/llamaindex_conversable_agent.py +0 -123
  77. autogen/agentchat/contrib/llava_agent.py +0 -176
  78. autogen/agentchat/contrib/math_user_proxy_agent.py +0 -471
  79. autogen/agentchat/contrib/multimodal_conversable_agent.py +0 -128
  80. autogen/agentchat/contrib/qdrant_retrieve_user_proxy_agent.py +0 -325
  81. autogen/agentchat/contrib/retrieve_assistant_agent.py +0 -56
  82. autogen/agentchat/contrib/retrieve_user_proxy_agent.py +0 -705
  83. autogen/agentchat/contrib/society_of_mind_agent.py +0 -203
  84. autogen/agentchat/contrib/swarm_agent.py +0 -463
  85. autogen/agentchat/contrib/text_analyzer_agent.py +0 -76
  86. autogen/agentchat/contrib/tool_retriever.py +0 -120
  87. autogen/agentchat/contrib/vectordb/__init__.py +0 -0
  88. autogen/agentchat/contrib/vectordb/base.py +0 -243
  89. autogen/agentchat/contrib/vectordb/chromadb.py +0 -326
  90. autogen/agentchat/contrib/vectordb/mongodb.py +0 -559
  91. autogen/agentchat/contrib/vectordb/pgvectordb.py +0 -958
  92. autogen/agentchat/contrib/vectordb/qdrant.py +0 -334
  93. autogen/agentchat/contrib/vectordb/utils.py +0 -126
  94. autogen/agentchat/contrib/web_surfer.py +0 -305
  95. autogen/agentchat/conversable_agent.py +0 -2908
  96. autogen/agentchat/groupchat.py +0 -1668
  97. autogen/agentchat/user_proxy_agent.py +0 -109
  98. autogen/agentchat/utils.py +0 -207
  99. autogen/browser_utils.py +0 -291
  100. autogen/cache/__init__.py +0 -10
  101. autogen/cache/abstract_cache_base.py +0 -78
  102. autogen/cache/cache.py +0 -182
  103. autogen/cache/cache_factory.py +0 -85
  104. autogen/cache/cosmos_db_cache.py +0 -150
  105. autogen/cache/disk_cache.py +0 -109
  106. autogen/cache/in_memory_cache.py +0 -61
  107. autogen/cache/redis_cache.py +0 -128
  108. autogen/code_utils.py +0 -745
  109. autogen/coding/__init__.py +0 -22
  110. autogen/coding/base.py +0 -113
  111. autogen/coding/docker_commandline_code_executor.py +0 -262
  112. autogen/coding/factory.py +0 -45
  113. autogen/coding/func_with_reqs.py +0 -203
  114. autogen/coding/jupyter/__init__.py +0 -22
  115. autogen/coding/jupyter/base.py +0 -32
  116. autogen/coding/jupyter/docker_jupyter_server.py +0 -164
  117. autogen/coding/jupyter/embedded_ipython_code_executor.py +0 -182
  118. autogen/coding/jupyter/jupyter_client.py +0 -224
  119. autogen/coding/jupyter/jupyter_code_executor.py +0 -161
  120. autogen/coding/jupyter/local_jupyter_server.py +0 -168
  121. autogen/coding/local_commandline_code_executor.py +0 -410
  122. autogen/coding/markdown_code_extractor.py +0 -44
  123. autogen/coding/utils.py +0 -57
  124. autogen/exception_utils.py +0 -46
  125. autogen/extensions/__init__.py +0 -0
  126. autogen/formatting_utils.py +0 -76
  127. autogen/function_utils.py +0 -362
  128. autogen/graph_utils.py +0 -148
  129. autogen/io/__init__.py +0 -15
  130. autogen/io/base.py +0 -105
  131. autogen/io/console.py +0 -43
  132. autogen/io/websockets.py +0 -213
  133. autogen/logger/__init__.py +0 -11
  134. autogen/logger/base_logger.py +0 -140
  135. autogen/logger/file_logger.py +0 -287
  136. autogen/logger/logger_factory.py +0 -29
  137. autogen/logger/logger_utils.py +0 -42
  138. autogen/logger/sqlite_logger.py +0 -459
  139. autogen/math_utils.py +0 -356
  140. autogen/oai/__init__.py +0 -33
  141. autogen/oai/anthropic.py +0 -428
  142. autogen/oai/bedrock.py +0 -606
  143. autogen/oai/cerebras.py +0 -270
  144. autogen/oai/client.py +0 -1148
  145. autogen/oai/client_utils.py +0 -167
  146. autogen/oai/cohere.py +0 -453
  147. autogen/oai/completion.py +0 -1216
  148. autogen/oai/gemini.py +0 -469
  149. autogen/oai/groq.py +0 -281
  150. autogen/oai/mistral.py +0 -279
  151. autogen/oai/ollama.py +0 -582
  152. autogen/oai/openai_utils.py +0 -811
  153. autogen/oai/together.py +0 -343
  154. autogen/retrieve_utils.py +0 -487
  155. autogen/runtime_logging.py +0 -163
  156. autogen/token_count_utils.py +0 -259
  157. autogen/types.py +0 -20
  158. autogen/version.py +0 -7
  159. {ag2-0.4.1.dist-info → ag2-0.4.2b1.dist-info}/LICENSE +0 -0
  160. {ag2-0.4.1.dist-info → ag2-0.4.2b1.dist-info}/NOTICE.md +0 -0
  161. {ag2-0.4.1.dist-info → ag2-0.4.2b1.dist-info}/WHEEL +0 -0
@@ -1,76 +0,0 @@
1
- # Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
2
- #
3
- # SPDX-License-Identifier: Apache-2.0
4
- #
5
- # Portions derived from https://github.com/microsoft/autogen are under the MIT License.
6
- # SPDX-License-Identifier: MIT
7
- from typing import Any, Dict, List, Literal, Optional, Tuple, Union
8
-
9
- from autogen.agentchat.agent import Agent
10
- from autogen.agentchat.assistant_agent import ConversableAgent
11
-
12
- system_message = """You are an expert in text analysis.
13
- The user will give you TEXT to analyze.
14
- The user will give you analysis INSTRUCTIONS copied twice, at both the beginning and the end.
15
- You will follow these INSTRUCTIONS in analyzing the TEXT, then give the results of your expert analysis in the format requested."""
16
-
17
-
18
- class TextAnalyzerAgent(ConversableAgent):
19
- """(Experimental) Text Analysis agent, a subclass of ConversableAgent designed to analyze text as instructed."""
20
-
21
- def __init__(
22
- self,
23
- name="analyzer",
24
- system_message: Optional[str] = system_message,
25
- human_input_mode: Literal["ALWAYS", "NEVER", "TERMINATE"] = "NEVER",
26
- llm_config: Optional[Union[Dict, bool]] = None,
27
- **kwargs,
28
- ):
29
- """
30
- Args:
31
- name (str): name of the agent.
32
- system_message (str): system message for the ChatCompletion inference.
33
- human_input_mode (str): This agent should NEVER prompt the human for input.
34
- llm_config (dict or False): llm inference configuration.
35
- Please refer to [OpenAIWrapper.create](/docs/reference/oai/client#create)
36
- for available options.
37
- To disable llm-based auto reply, set to False.
38
- **kwargs (dict): other kwargs in [ConversableAgent](../conversable_agent#__init__).
39
- """
40
- super().__init__(
41
- name=name,
42
- system_message=system_message,
43
- human_input_mode=human_input_mode,
44
- llm_config=llm_config,
45
- **kwargs,
46
- )
47
- self.register_reply(Agent, TextAnalyzerAgent._analyze_in_reply, position=2)
48
-
49
- def _analyze_in_reply(
50
- self,
51
- messages: Optional[List[Dict]] = None,
52
- sender: Optional[Agent] = None,
53
- config: Optional[Any] = None,
54
- ) -> Tuple[bool, Union[str, Dict, None]]:
55
- """Analyzes the given text as instructed, and returns the analysis as a message.
56
- Assumes exactly two messages containing the text to analyze and the analysis instructions.
57
- See Teachability.analyze for an example of how to use this method."""
58
- if self.llm_config is False:
59
- raise ValueError("TextAnalyzerAgent requires self.llm_config to be set in its base class.")
60
- if messages is None:
61
- messages = self._oai_messages[sender] # In case of a direct call.
62
- assert len(messages) == 2
63
-
64
- # Delegate to the analysis method.
65
- return True, self.analyze_text(messages[0]["content"], messages[1]["content"])
66
-
67
- def analyze_text(self, text_to_analyze, analysis_instructions):
68
- """Analyzes the given text as instructed, and returns the analysis."""
69
- # Assemble the message.
70
- text_to_analyze = "# TEXT\n" + text_to_analyze + "\n"
71
- analysis_instructions = "# INSTRUCTIONS\n" + analysis_instructions + "\n"
72
- msg_text = "\n".join(
73
- [analysis_instructions, text_to_analyze, analysis_instructions]
74
- ) # Repeat the instructions.
75
- # Generate and return the analysis string.
76
- return self.generate_oai_reply([{"role": "user", "content": msg_text}], None, None)[1]
@@ -1,120 +0,0 @@
1
- # Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
2
- #
3
- # SPDX-License-Identifier: Apache-2.0
4
- #
5
- # Portions derived from https://github.com/microsoft/autogen are under the MIT License.
6
- # SPDX-License-Identifier: MIT
7
- import importlib.util
8
- import inspect
9
- import os
10
- from textwrap import dedent, indent
11
-
12
- import pandas as pd
13
- from sentence_transformers import SentenceTransformer, util
14
-
15
- from autogen import AssistantAgent, UserProxyAgent
16
- from autogen.coding import LocalCommandLineCodeExecutor
17
-
18
-
19
- class ToolBuilder:
20
- TOOL_USING_PROMPT = """# Functions
21
- You have access to the following functions. They can be accessed from the module called 'functions' by their function names.
22
- For example, if there is a function called `foo` you could import it by writing `from functions import foo`
23
- {functions}
24
- """
25
-
26
- def __init__(self, corpus_path, retriever="all-mpnet-base-v2"):
27
-
28
- self.df = pd.read_csv(corpus_path, sep="\t")
29
- document_list = self.df["document_content"].tolist()
30
-
31
- self.model = SentenceTransformer(retriever)
32
- self.embeddings = self.model.encode(document_list)
33
-
34
- def retrieve(self, query, top_k=3):
35
- # Encode the query using the Sentence Transformer model
36
- query_embedding = self.model.encode([query])
37
-
38
- hits = util.semantic_search(query_embedding, self.embeddings, top_k=top_k)
39
-
40
- results = []
41
- for hit in hits[0]:
42
- results.append(self.df.iloc[hit["corpus_id"], 1])
43
- return results
44
-
45
- def bind(self, agent: AssistantAgent, functions: str):
46
- """Binds the function to the agent so that agent is aware of it."""
47
- sys_message = agent.system_message
48
- sys_message += self.TOOL_USING_PROMPT.format(functions=functions)
49
- agent.update_system_message(sys_message)
50
- return
51
-
52
- def bind_user_proxy(self, agent: UserProxyAgent, tool_root: str):
53
- """
54
- Updates user proxy agent with a executor so that code executor can successfully execute function-related code.
55
- Returns an updated user proxy.
56
- """
57
- # Find all the functions in the tool root
58
- functions = find_callables(tool_root)
59
-
60
- code_execution_config = agent._code_execution_config
61
- executor = LocalCommandLineCodeExecutor(
62
- timeout=code_execution_config.get("timeout", 180),
63
- work_dir=code_execution_config.get("work_dir", "coding"),
64
- functions=functions,
65
- )
66
- code_execution_config = {
67
- "executor": executor,
68
- "last_n_messages": code_execution_config.get("last_n_messages", 1),
69
- }
70
- updated_user_proxy = UserProxyAgent(
71
- name=agent.name,
72
- is_termination_msg=agent._is_termination_msg,
73
- code_execution_config=code_execution_config,
74
- human_input_mode="NEVER",
75
- default_auto_reply=agent._default_auto_reply,
76
- )
77
- return updated_user_proxy
78
-
79
-
80
- def get_full_tool_description(py_file):
81
- """
82
- Retrieves the function signature for a given Python file.
83
- """
84
- with open(py_file, "r") as f:
85
- code = f.read()
86
- exec(code)
87
- function_name = os.path.splitext(os.path.basename(py_file))[0]
88
- if function_name in locals():
89
- func = locals()[function_name]
90
- content = f"def {func.__name__}{inspect.signature(func)}:\n"
91
- docstring = func.__doc__
92
-
93
- if docstring:
94
- docstring = dedent(docstring)
95
- docstring = '"""' + docstring + '"""'
96
- docstring = indent(docstring, " ")
97
- content += docstring + "\n"
98
- return content
99
- else:
100
- raise ValueError(f"Function {function_name} not found in {py_file}")
101
-
102
-
103
- def find_callables(directory):
104
- """
105
- Find all callable objects defined in Python files within the specified directory.
106
- """
107
- callables = []
108
- for root, dirs, files in os.walk(directory):
109
- for file in files:
110
- if file.endswith(".py"):
111
- module_name = os.path.splitext(file)[0]
112
- module_path = os.path.join(root, file)
113
- spec = importlib.util.spec_from_file_location(module_name, module_path)
114
- module = importlib.util.module_from_spec(spec)
115
- spec.loader.exec_module(module)
116
- for name, value in module.__dict__.items():
117
- if callable(value) and name == module_name:
118
- callables.append(value)
119
- break
120
- return callables
File without changes
@@ -1,243 +0,0 @@
1
- # Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai
2
- #
3
- # SPDX-License-Identifier: Apache-2.0
4
- #
5
- # Portions derived from https://github.com/microsoft/autogen are under the MIT License.
6
- # SPDX-License-Identifier: MIT
7
- from typing import (
8
- Any,
9
- Callable,
10
- List,
11
- Mapping,
12
- Optional,
13
- Protocol,
14
- Sequence,
15
- Tuple,
16
- TypedDict,
17
- Union,
18
- runtime_checkable,
19
- )
20
-
21
- Metadata = Union[Mapping[str, Any], None]
22
- Vector = Union[Sequence[float], Sequence[int]]
23
- ItemID = Union[str, int] # chromadb doesn't support int ids, VikingDB does
24
-
25
-
26
- class Document(TypedDict):
27
- """A Document is a record in the vector database.
28
-
29
- id: ItemID | the unique identifier of the document.
30
- content: str | the text content of the chunk.
31
- metadata: Metadata, Optional | contains additional information about the document such as source, date, etc.
32
- embedding: Vector, Optional | the vector representation of the content.
33
- """
34
-
35
- id: ItemID
36
- content: str
37
- metadata: Optional[Metadata]
38
- embedding: Optional[Vector]
39
-
40
-
41
- """QueryResults is the response from the vector database for a query/queries.
42
- A query is a list containing one string while queries is a list containing multiple strings.
43
- The response is a list of query results, each query result is a list of tuples containing the document and the distance.
44
- """
45
- QueryResults = List[List[Tuple[Document, float]]]
46
-
47
-
48
- @runtime_checkable
49
- class VectorDB(Protocol):
50
- """
51
- Abstract class for vector database. A vector database is responsible for storing and retrieving documents.
52
-
53
- Attributes:
54
- active_collection: Any | The active collection in the vector database. Make get_collection faster. Default is None.
55
- type: str | The type of the vector database, chroma, pgvector, etc. Default is "".
56
-
57
- Methods:
58
- create_collection: Callable[[str, bool, bool], Any] | Create a collection in the vector database.
59
- get_collection: Callable[[str], Any] | Get the collection from the vector database.
60
- delete_collection: Callable[[str], Any] | Delete the collection from the vector database.
61
- insert_docs: Callable[[List[Document], str, bool], None] | Insert documents into the collection of the vector database.
62
- update_docs: Callable[[List[Document], str], None] | Update documents in the collection of the vector database.
63
- delete_docs: Callable[[List[ItemID], str], None] | Delete documents from the collection of the vector database.
64
- retrieve_docs: Callable[[List[str], str, int, float], QueryResults] | Retrieve documents from the collection of the vector database based on the queries.
65
- get_docs_by_ids: Callable[[List[ItemID], str], List[Document]] | Retrieve documents from the collection of the vector database based on the ids.
66
- """
67
-
68
- active_collection: Any = None
69
- type: str = ""
70
- embedding_function: Optional[Callable[[List[str]], List[List[float]]]] = (
71
- None # embeddings = embedding_function(sentences)
72
- )
73
-
74
- def create_collection(self, collection_name: str, overwrite: bool = False, get_or_create: bool = True) -> Any:
75
- """
76
- Create a collection in the vector database.
77
- Case 1. if the collection does not exist, create the collection.
78
- Case 2. the collection exists, if overwrite is True, it will overwrite the collection.
79
- Case 3. the collection exists and overwrite is False, if get_or_create is True, it will get the collection,
80
- otherwise it raise a ValueError.
81
-
82
- Args:
83
- collection_name: str | The name of the collection.
84
- overwrite: bool | Whether to overwrite the collection if it exists. Default is False.
85
- get_or_create: bool | Whether to get the collection if it exists. Default is True.
86
-
87
- Returns:
88
- Any | The collection object.
89
- """
90
- ...
91
-
92
- def get_collection(self, collection_name: str = None) -> Any:
93
- """
94
- Get the collection from the vector database.
95
-
96
- Args:
97
- collection_name: str | The name of the collection. Default is None. If None, return the
98
- current active collection.
99
-
100
- Returns:
101
- Any | The collection object.
102
- """
103
- ...
104
-
105
- def delete_collection(self, collection_name: str) -> Any:
106
- """
107
- Delete the collection from the vector database.
108
-
109
- Args:
110
- collection_name: str | The name of the collection.
111
-
112
- Returns:
113
- Any
114
- """
115
- ...
116
-
117
- def insert_docs(self, docs: List[Document], collection_name: str = None, upsert: bool = False, **kwargs) -> None:
118
- """
119
- Insert documents into the collection of the vector database.
120
-
121
- Args:
122
- docs: List[Document] | A list of documents. Each document is a TypedDict `Document`.
123
- collection_name: str | The name of the collection. Default is None.
124
- upsert: bool | Whether to update the document if it exists. Default is False.
125
- kwargs: Dict | Additional keyword arguments.
126
-
127
- Returns:
128
- None
129
- """
130
- ...
131
-
132
- def update_docs(self, docs: List[Document], collection_name: str = None, **kwargs) -> None:
133
- """
134
- Update documents in the collection of the vector database.
135
-
136
- Args:
137
- docs: List[Document] | A list of documents.
138
- collection_name: str | The name of the collection. Default is None.
139
- kwargs: Dict | Additional keyword arguments.
140
-
141
- Returns:
142
- None
143
- """
144
- ...
145
-
146
- def delete_docs(self, ids: List[ItemID], collection_name: str = None, **kwargs) -> None:
147
- """
148
- Delete documents from the collection of the vector database.
149
-
150
- Args:
151
- ids: List[ItemID] | A list of document ids. Each id is a typed `ItemID`.
152
- collection_name: str | The name of the collection. Default is None.
153
- kwargs: Dict | Additional keyword arguments.
154
-
155
- Returns:
156
- None
157
- """
158
- ...
159
-
160
- def retrieve_docs(
161
- self,
162
- queries: List[str],
163
- collection_name: str = None,
164
- n_results: int = 10,
165
- distance_threshold: float = -1,
166
- **kwargs,
167
- ) -> QueryResults:
168
- """
169
- Retrieve documents from the collection of the vector database based on the queries.
170
-
171
- Args:
172
- queries: List[str] | A list of queries. Each query is a string.
173
- collection_name: str | The name of the collection. Default is None.
174
- n_results: int | The number of relevant documents to return. Default is 10.
175
- distance_threshold: float | The threshold for the distance score, only distance smaller than it will be
176
- returned. Don't filter with it if < 0. Default is -1.
177
- kwargs: Dict | Additional keyword arguments.
178
-
179
- Returns:
180
- QueryResults | The query results. Each query result is a list of list of tuples containing the document and
181
- the distance.
182
- """
183
- ...
184
-
185
- def get_docs_by_ids(
186
- self, ids: List[ItemID] = None, collection_name: str = None, include=None, **kwargs
187
- ) -> List[Document]:
188
- """
189
- Retrieve documents from the collection of the vector database based on the ids.
190
-
191
- Args:
192
- ids: List[ItemID] | A list of document ids. If None, will return all the documents. Default is None.
193
- collection_name: str | The name of the collection. Default is None.
194
- include: List[str] | The fields to include. Default is None.
195
- If None, will include ["metadatas", "documents"], ids will always be included. This may differ
196
- depending on the implementation.
197
- kwargs: dict | Additional keyword arguments.
198
-
199
- Returns:
200
- List[Document] | The results.
201
- """
202
- ...
203
-
204
-
205
- class VectorDBFactory:
206
- """
207
- Factory class for creating vector databases.
208
- """
209
-
210
- PREDEFINED_VECTOR_DB = ["chroma", "pgvector", "mongodb", "qdrant"]
211
-
212
- @staticmethod
213
- def create_vector_db(db_type: str, **kwargs) -> VectorDB:
214
- """
215
- Create a vector database.
216
-
217
- Args:
218
- db_type: str | The type of the vector database.
219
- kwargs: Dict | The keyword arguments for initializing the vector database.
220
-
221
- Returns:
222
- VectorDB | The vector database.
223
- """
224
- if db_type.lower() in ["chroma", "chromadb"]:
225
- from .chromadb import ChromaVectorDB
226
-
227
- return ChromaVectorDB(**kwargs)
228
- if db_type.lower() in ["pgvector", "pgvectordb"]:
229
- from .pgvectordb import PGVectorDB
230
-
231
- return PGVectorDB(**kwargs)
232
- if db_type.lower() in ["mdb", "mongodb", "atlas"]:
233
- from .mongodb import MongoDBAtlasVectorDB
234
-
235
- return MongoDBAtlasVectorDB(**kwargs)
236
- if db_type.lower() in ["qdrant", "qdrantdb"]:
237
- from .qdrant import QdrantVectorDB
238
-
239
- return QdrantVectorDB(**kwargs)
240
- else:
241
- raise ValueError(
242
- f"Unsupported vector database type: {db_type}. Valid types are {VectorDBFactory.PREDEFINED_VECTOR_DB}."
243
- )