langroid 0.20.0__tar.gz → 0.21.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.
Files changed (149) hide show
  1. {langroid-0.20.0 → langroid-0.21.0}/PKG-INFO +4 -1
  2. {langroid-0.20.0 → langroid-0.21.0}/README.md +3 -0
  3. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/base.py +10 -1
  4. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/arangodb/arangodb_agent.py +236 -69
  5. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/arangodb/system_messages.py +39 -10
  6. langroid-0.21.0/langroid/agent/special/arangodb/tools.py +102 -0
  7. langroid-0.21.0/langroid/agent/special/arangodb/utils.py +36 -0
  8. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/task.py +77 -55
  9. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/orchestration.py +7 -7
  10. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/__init__.py +2 -0
  11. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/openai_gpt.py +24 -6
  12. {langroid-0.20.0 → langroid-0.21.0}/pyproject.toml +1 -1
  13. langroid-0.20.0/langroid/agent/special/arangodb/tools.py +0 -39
  14. {langroid-0.20.0 → langroid-0.21.0}/LICENSE +0 -0
  15. {langroid-0.20.0 → langroid-0.21.0}/langroid/__init__.py +0 -0
  16. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/__init__.py +0 -0
  17. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/batch.py +0 -0
  18. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/callbacks/__init__.py +0 -0
  19. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/callbacks/chainlit.py +0 -0
  20. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/chat_agent.py +0 -0
  21. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/chat_document.py +0 -0
  22. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/helpers.py +0 -0
  23. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/junk +0 -0
  24. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/openai_assistant.py +0 -0
  25. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/__init__.py +0 -0
  26. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/arangodb/__init__.py +0 -0
  27. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/doc_chat_agent.py +0 -0
  28. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/lance_doc_chat_agent.py +0 -0
  29. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/lance_rag/__init__.py +0 -0
  30. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/lance_rag/critic_agent.py +0 -0
  31. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/lance_rag/lance_rag_task.py +0 -0
  32. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/lance_rag/query_planner_agent.py +0 -0
  33. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/lance_tools.py +0 -0
  34. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/neo4j/__init__.py +0 -0
  35. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/neo4j/csv_kg_chat.py +0 -0
  36. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/neo4j/neo4j_chat_agent.py +0 -0
  37. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/neo4j/system_messages.py +0 -0
  38. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/neo4j/tools.py +0 -0
  39. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/relevance_extractor_agent.py +0 -0
  40. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/retriever_agent.py +0 -0
  41. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/sql/__init__.py +0 -0
  42. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/sql/sql_chat_agent.py +0 -0
  43. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/sql/utils/__init__.py +0 -0
  44. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/sql/utils/description_extractors.py +0 -0
  45. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/sql/utils/populate_metadata.py +0 -0
  46. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/sql/utils/system_message.py +0 -0
  47. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/sql/utils/tools.py +0 -0
  48. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/special/table_chat_agent.py +0 -0
  49. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/structured_message.py +0 -0
  50. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tool_message.py +0 -0
  51. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/__init__.py +0 -0
  52. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/duckduckgo_search_tool.py +0 -0
  53. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/file_tools.py +0 -0
  54. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/google_search_tool.py +0 -0
  55. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/metaphor_search_tool.py +0 -0
  56. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/recipient_tool.py +0 -0
  57. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/retrieval_tool.py +0 -0
  58. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/rewind_tool.py +0 -0
  59. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/tools/segment_extract_tool.py +0 -0
  60. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/typed_task.py +0 -0
  61. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent/xml_tool_message.py +0 -0
  62. {langroid-0.20.0 → langroid-0.21.0}/langroid/agent_config.py +0 -0
  63. {langroid-0.20.0 → langroid-0.21.0}/langroid/cachedb/__init__.py +0 -0
  64. {langroid-0.20.0 → langroid-0.21.0}/langroid/cachedb/base.py +0 -0
  65. {langroid-0.20.0 → langroid-0.21.0}/langroid/cachedb/momento_cachedb.py +0 -0
  66. {langroid-0.20.0 → langroid-0.21.0}/langroid/cachedb/redis_cachedb.py +0 -0
  67. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/__init__.py +0 -0
  68. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/base.py +0 -0
  69. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/clustering.py +0 -0
  70. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/models.py +0 -0
  71. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/protoc/__init__.py +0 -0
  72. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/protoc/embeddings.proto +0 -0
  73. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/protoc/embeddings_pb2.py +0 -0
  74. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/protoc/embeddings_pb2.pyi +0 -0
  75. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/protoc/embeddings_pb2_grpc.py +0 -0
  76. {langroid-0.20.0 → langroid-0.21.0}/langroid/embedding_models/remote_embeds.py +0 -0
  77. {langroid-0.20.0 → langroid-0.21.0}/langroid/exceptions.py +0 -0
  78. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/.chainlit/config.toml +0 -0
  79. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/.chainlit/translations/en-US.json +0 -0
  80. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/azure_openai.py +0 -0
  81. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/base.py +0 -0
  82. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/config.py +0 -0
  83. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/mock_lm.py +0 -0
  84. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/prompt_formatter/__init__.py +0 -0
  85. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/prompt_formatter/base.py +0 -0
  86. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/prompt_formatter/hf_formatter.py +0 -0
  87. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/prompt_formatter/llama2_formatter.py +0 -0
  88. {langroid-0.20.0 → langroid-0.21.0}/langroid/language_models/utils.py +0 -0
  89. {langroid-0.20.0 → langroid-0.21.0}/langroid/mytypes.py +0 -0
  90. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/__init__.py +0 -0
  91. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/agent_chats.py +0 -0
  92. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/code-parsing.md +0 -0
  93. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/code_parser.py +0 -0
  94. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/config.py +0 -0
  95. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/document_parser.py +0 -0
  96. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/image_text.py +0 -0
  97. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/para_sentence_split.py +0 -0
  98. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/parse_json.py +0 -0
  99. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/parser.py +0 -0
  100. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/repo_loader.py +0 -0
  101. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/routing.py +0 -0
  102. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/search.py +0 -0
  103. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/spider.py +0 -0
  104. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/table_loader.py +0 -0
  105. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/url_loader.py +0 -0
  106. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/url_loader_cookies.py +0 -0
  107. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/urls.py +0 -0
  108. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/utils.py +0 -0
  109. {langroid-0.20.0 → langroid-0.21.0}/langroid/parsing/web_search.py +0 -0
  110. {langroid-0.20.0 → langroid-0.21.0}/langroid/prompts/__init__.py +0 -0
  111. {langroid-0.20.0 → langroid-0.21.0}/langroid/prompts/chat-gpt4-system-prompt.md +0 -0
  112. {langroid-0.20.0 → langroid-0.21.0}/langroid/prompts/dialog.py +0 -0
  113. {langroid-0.20.0 → langroid-0.21.0}/langroid/prompts/prompts_config.py +0 -0
  114. {langroid-0.20.0 → langroid-0.21.0}/langroid/prompts/templates.py +0 -0
  115. {langroid-0.20.0 → langroid-0.21.0}/langroid/py.typed +0 -0
  116. {langroid-0.20.0 → langroid-0.21.0}/langroid/pydantic_v1/__init__.py +0 -0
  117. {langroid-0.20.0 → langroid-0.21.0}/langroid/pydantic_v1/main.py +0 -0
  118. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/.chainlit/config.toml +0 -0
  119. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/.chainlit/translations/en-US.json +0 -0
  120. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/__init__.py +0 -0
  121. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/algorithms/__init__.py +0 -0
  122. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/algorithms/graph.py +0 -0
  123. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/configuration.py +0 -0
  124. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/constants.py +0 -0
  125. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/docker.py +0 -0
  126. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/git_utils.py +0 -0
  127. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/globals.py +0 -0
  128. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/llms/__init__.py +0 -0
  129. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/llms/strings.py +0 -0
  130. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/logging.py +0 -0
  131. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/object_registry.py +0 -0
  132. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/output/__init__.py +0 -0
  133. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/output/citations.py +0 -0
  134. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/output/printing.py +0 -0
  135. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/output/status.py +0 -0
  136. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/pandas_utils.py +0 -0
  137. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/pydantic_utils.py +0 -0
  138. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/system.py +0 -0
  139. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/types.py +0 -0
  140. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/web/__init__.py +0 -0
  141. {langroid-0.20.0 → langroid-0.21.0}/langroid/utils/web/login.py +0 -0
  142. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/__init__.py +0 -0
  143. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/base.py +0 -0
  144. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/chromadb.py +0 -0
  145. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/lancedb.py +0 -0
  146. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/meilisearch.py +0 -0
  147. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/momento.py +0 -0
  148. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/qdrant_cloud.py +0 -0
  149. {langroid-0.20.0 → langroid-0.21.0}/langroid/vector_store/qdrantdb.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.1
2
2
  Name: langroid
3
- Version: 0.20.0
3
+ Version: 0.21.0
4
4
  Summary: Harness LLMs with Multi-Agent Programming
5
5
  License: MIT
6
6
  Author: Prasad Chalasani
@@ -248,6 +248,9 @@ teacher_task.run()
248
248
  <details>
249
249
  <summary> <b>Click to expand</b></summary>
250
250
 
251
+ - **Nov 2024:**
252
+ - **[0.20.0](https://github.com/langroid/langroid/releases/tag/0.20.0)** Support for
253
+ ArangoDB Knowledge Graphs.
251
254
  - **Oct 2024:**
252
255
  - **[0.18.0]** [LLMConfig.async_stream_quiet](https://langroid.github.io/langroid/notes/async-streaming/) flag to
253
256
  turn off LLM output in async + stream mode.
@@ -135,6 +135,9 @@ teacher_task.run()
135
135
  <details>
136
136
  <summary> <b>Click to expand</b></summary>
137
137
 
138
+ - **Nov 2024:**
139
+ - **[0.20.0](https://github.com/langroid/langroid/releases/tag/0.20.0)** Support for
140
+ ArangoDB Knowledge Graphs.
138
141
  - **Oct 2024:**
139
142
  - **[0.18.0]** [LLMConfig.async_stream_quiet](https://langroid.github.io/langroid/notes/async-streaming/) flag to
140
143
  turn off LLM output in async + stream mode.
@@ -880,7 +880,13 @@ class Agent(ABC):
880
880
  return cdoc
881
881
 
882
882
  def has_tool_message_attempt(self, msg: str | ChatDocument | None) -> bool:
883
- """Check whether msg contains a Tool/fn-call attempt (by the LLM)"""
883
+ """
884
+ Check whether msg contains a Tool/fn-call attempt (by the LLM).
885
+
886
+ CAUTION: This uses self.get_tool_messages(msg) which as a side-effect
887
+ may update msg.tool_messages when msg is a ChatDocument, if there are
888
+ any tools in msg.
889
+ """
884
890
  if msg is None:
885
891
  return False
886
892
  try:
@@ -921,6 +927,9 @@ class Agent(ABC):
921
927
  ) -> List[ToolMessage]:
922
928
  """
923
929
  Get ToolMessages recognized in msg, handle-able by this agent.
930
+ NOTE: as a side-effect, this will update msg.tool_messages
931
+ when msg is a ChatDocument and msg contains tool messages.
932
+
924
933
  If all_tools is True:
925
934
  - return all tools, i.e. any tool in self.llm_tools_known,
926
935
  whether it is handled by this agent or not;
@@ -27,6 +27,7 @@ from langroid.agent.special.arangodb.tools import (
27
27
  aql_retrieval_tool_name,
28
28
  arango_schema_tool_name,
29
29
  )
30
+ from langroid.agent.special.arangodb.utils import count_fields, trim_schema
30
31
  from langroid.agent.tools.orchestration import DoneTool, ForwardTool
31
32
  from langroid.exceptions import LangroidImportError
32
33
  from langroid.mytypes import Entity
@@ -88,11 +89,14 @@ class QueryResult(BaseModel):
88
89
  class ArangoChatAgentConfig(ChatAgentConfig):
89
90
  arango_settings: ArangoSettings = ArangoSettings()
90
91
  system_message: str = DEFAULT_ARANGO_CHAT_SYSTEM_MESSAGE
91
- kg_schema: Optional[Dict[str, List[Dict[str, Any]]]] = None
92
+ kg_schema: str | Dict[str, List[Dict[str, Any]]] | None = None
92
93
  database_created: bool = False
93
- use_schema_tools: bool = True
94
+ prepopulate_schema: bool = True
94
95
  use_functions_api: bool = True
96
+ max_num_results: int = 10 # how many results to return from AQL query
95
97
  max_result_tokens: int = 1000 # truncate long results to this many tokens
98
+ max_schema_fields: int = 500 # max fields to show in schema
99
+ max_tries: int = 10 # how many attempts to answer user question
96
100
  use_tools: bool = False
97
101
  schema_sample_pct: float = 0
98
102
  # whether the agent is used in a continuous chat with user,
@@ -103,16 +107,81 @@ class ArangoChatAgentConfig(ChatAgentConfig):
103
107
 
104
108
  class ArangoChatAgent(ChatAgent):
105
109
  def __init__(self, config: ArangoChatAgentConfig):
110
+ super().__init__(config)
106
111
  self.config: ArangoChatAgentConfig = config
112
+ self.init_state()
107
113
  self._validate_config()
108
114
  self._import_arango()
109
115
  self._initialize_db()
110
116
  self._init_tools_sys_message()
111
- self.init_state()
112
117
 
113
118
  def init_state(self) -> None:
114
119
  super().init_state()
115
120
  self.current_retrieval_aql_query: str = ""
121
+ self.current_schema_params: ArangoSchemaTool = ArangoSchemaTool()
122
+ self.num_tries = 0 # how many attempts to answer user question
123
+
124
+ def user_response(
125
+ self,
126
+ msg: Optional[str | ChatDocument] = None,
127
+ ) -> Optional[ChatDocument]:
128
+ response = super().user_response(msg)
129
+ if response is None:
130
+ return None
131
+ response_str = response.content if response is not None else ""
132
+ if response_str != "":
133
+ self.num_tries = 0 # reset number of tries if user responds
134
+ return response
135
+
136
+ def llm_response(
137
+ self, message: Optional[str | ChatDocument] = None
138
+ ) -> Optional[ChatDocument]:
139
+ if self.num_tries > self.config.max_tries:
140
+ if self.config.chat_mode:
141
+ return self.create_llm_response(
142
+ content=f"""
143
+ {self.config.addressing_prefix}User
144
+ I give up, since I have exceeded the
145
+ maximum number of tries ({self.config.max_tries}).
146
+ Feel free to give me some hints!
147
+ """
148
+ )
149
+ else:
150
+ return self.create_llm_response(
151
+ tool_messages=[
152
+ DoneTool(
153
+ content=f"""
154
+ Exceeded maximum number of tries ({self.config.max_tries}).
155
+ """
156
+ )
157
+ ]
158
+ )
159
+
160
+ if isinstance(message, ChatDocument) and message.metadata.sender == Entity.USER:
161
+ message.content = (
162
+ message.content
163
+ + "\n"
164
+ + """
165
+ (REMEMBER, Do NOT use more than ONE TOOL/FUNCTION at a time!
166
+ you must WAIT for a helper to send you the RESULT(S) before
167
+ making another TOOL/FUNCTION call)
168
+ """
169
+ )
170
+
171
+ response = super().llm_response(message)
172
+ if (
173
+ response is not None
174
+ and self.config.chat_mode
175
+ and self.config.addressing_prefix in response.content
176
+ and self.has_tool_message_attempt(response)
177
+ ):
178
+ # response contains both a user-addressing and a tool, which
179
+ # is not allowed, so remove the user-addressing prefix
180
+ response.content = response.content.replace(
181
+ self.config.addressing_prefix, ""
182
+ )
183
+
184
+ return response
116
185
 
117
186
  def _validate_config(self) -> None:
118
187
  assert isinstance(self.config, ArangoChatAgentConfig)
@@ -230,6 +299,7 @@ class ArangoChatAgent(ChatAgent):
230
299
  try:
231
300
  cursor = self.db.aql.execute(query, bind_vars=bind_vars)
232
301
  records = [doc for doc in cursor] # type: ignore
302
+ records = records[: self.config.max_num_results]
233
303
  logger.warning(f"Records retrieved: {records}")
234
304
  return QueryResult(success=True, data=records if records else [])
235
305
  except Exception as e:
@@ -273,6 +343,28 @@ class ArangoChatAgent(ChatAgent):
273
343
  success=False, data=f"Failed after max retries: {str(e)}"
274
344
  )
275
345
 
346
+ def _limit_tokens(self, text: str) -> str:
347
+ result = text
348
+ n_toks = self.num_tokens(result)
349
+ if n_toks > self.config.max_result_tokens:
350
+ logger.warning(
351
+ f"""
352
+ Your query resulted in a large result of
353
+ {n_toks} tokens,
354
+ which will be truncated to {self.config.max_result_tokens} tokens.
355
+ If this does not give satisfactory results,
356
+ please retry with a more focused query.
357
+ """
358
+ )
359
+ if self.parser is not None:
360
+ result = self.parser.truncate_tokens(
361
+ result,
362
+ self.config.max_result_tokens,
363
+ )
364
+ else:
365
+ result = result[: self.config.max_result_tokens * 4] # truncate roughly
366
+ return result
367
+
276
368
  def aql_retrieval_tool(self, msg: AQLRetrievalTool) -> str:
277
369
  """Handle AQL query for data retrieval"""
278
370
  if not self.tried_schema:
@@ -285,7 +377,13 @@ class ArangoChatAgent(ChatAgent):
285
377
  return """
286
378
  You need to create the database first using `{aql_creation_tool_name}`.
287
379
  """
380
+ self.num_tries += 1
288
381
  query = msg.aql_query
382
+ if query == self.current_retrieval_aql_query:
383
+ return """
384
+ You have already tried this query, so you will get the same results again!
385
+ If you need to retry, please MODIFY the query to get different results.
386
+ """
289
387
  self.current_retrieval_aql_query = query
290
388
  logger.info(f"Executing AQL query: {query}")
291
389
  response = self.read_query(query)
@@ -299,28 +397,11 @@ class ArangoChatAgent(ChatAgent):
299
397
  """
300
398
  # truncate long results
301
399
  result = str(response.data)
302
- n_toks = self.num_tokens(result)
303
- if n_toks > self.config.max_result_tokens:
304
- logger.warning(
305
- f"""
306
- Your query resulted in a large result of
307
- {n_toks} tokens,
308
- which will be truncated to {self.config.max_result_tokens} tokens.
309
- If this does not give satisfactory results,
310
- please retry with a more focused query.
311
- """
312
- )
313
- if self.parser is not None:
314
- result = self.parser.truncate_tokens(
315
- result,
316
- self.config.max_result_tokens,
317
- )
318
- else:
319
- result = result[: self.config.max_result_tokens * 4] # truncate roughly
320
- return result
400
+ return self._limit_tokens(result)
321
401
 
322
402
  def aql_creation_tool(self, msg: AQLCreationTool) -> str:
323
403
  """Handle AQL query for creating data"""
404
+ self.num_tries += 1
324
405
  query = msg.aql_query
325
406
  logger.info(f"Executing AQL query: {query}")
326
407
  response = self.write_query(query)
@@ -334,12 +415,45 @@ class ArangoChatAgent(ChatAgent):
334
415
  self,
335
416
  msg: ArangoSchemaTool | None,
336
417
  ) -> Dict[str, List[Dict[str, Any]]] | str:
337
- """Get database schema including collections, properties, and relationships"""
418
+ """Get database schema. If collections=None, include all collections.
419
+ If properties=False, show only connection info,
420
+ else show all properties and example-docs.
421
+ """
422
+
423
+ if (
424
+ msg is not None
425
+ and msg.collections == self.current_schema_params.collections
426
+ and msg.properties == self.current_schema_params.properties
427
+ ):
428
+ return """
429
+ You have already tried this schema TOOL, so you will get the same results
430
+ again! Please MODIFY the tool params `collections` or `properties` to get
431
+ different results.
432
+ """
433
+
434
+ if msg is not None:
435
+ collections = msg.collections
436
+ properties = msg.properties
437
+ else:
438
+ collections = None
439
+ properties = True
338
440
  self.tried_schema = True
339
- if self.config.kg_schema is not None and len(self.config.kg_schema) > 0:
441
+ if (
442
+ self.config.kg_schema is not None
443
+ and len(self.config.kg_schema) > 0
444
+ and msg is None
445
+ ):
446
+ # we are trying to pre-populate full schema before the agent runs,
447
+ # so get it if it's already available
448
+ # (Note of course that this "full schema" may actually be incomplete)
340
449
  return self.config.kg_schema
450
+
451
+ # increment tries only if the LLM is asking for the schema,
452
+ # in which case msg will not be None
453
+ self.num_tries += msg is not None
454
+
341
455
  try:
342
- # Get graph schemas
456
+ # Get graph schemas (keeping full graph info)
343
457
  graph_schema = [
344
458
  {"graph_name": g["name"], "edge_definitions": g["edge_definitions"]}
345
459
  for g in self.db.graphs() # type: ignore
@@ -348,57 +462,78 @@ class ArangoChatAgent(ChatAgent):
348
462
  # Get collection schemas
349
463
  collection_schema = []
350
464
  for collection in self.db.collections(): # type: ignore
351
- if collection["name"].startswith("_"): # Skip system collections
465
+ if collection["name"].startswith("_"):
352
466
  continue
353
467
 
354
468
  col_name = collection["name"]
469
+ if collections and col_name not in collections:
470
+ continue
471
+
355
472
  col_type = collection["type"]
356
473
  col_size = self.db.collection(col_name).count()
357
474
 
358
- if col_size == 0: # Skip empty collections
475
+ if col_size == 0:
359
476
  continue
360
477
 
361
- # Calculate sample size
362
- limit_amount = (
363
- ceil(
364
- self.config.schema_sample_pct * col_size / 100.0 # type: ignore
365
- )
366
- or 1
367
- )
368
-
369
- # Query to get sample documents and their properties
370
- sample_query = f"""
371
- FOR doc in {col_name}
372
- LIMIT {limit_amount}
373
- RETURN doc
374
- """
478
+ if properties:
479
+ # Full property collection with sampling
480
+ lim = self.config.schema_sample_pct * col_size # type: ignore
481
+ limit_amount = ceil(lim / 100.0) or 1
482
+ sample_query = f"""
483
+ FOR doc in {col_name}
484
+ LIMIT {limit_amount}
485
+ RETURN doc
486
+ """
375
487
 
376
- properties = []
377
- example_doc = None
378
-
379
- def simplify_doc(doc: Any) -> Any:
380
- if isinstance(doc, list) and len(doc) > 0:
381
- return [simplify_doc(doc[0])]
382
- if isinstance(doc, dict):
383
- return {k: simplify_doc(v) for k, v in doc.items()}
384
- return doc
385
-
386
- for doc in self.db.aql.execute(sample_query): # type: ignore
387
- if example_doc is None:
388
- example_doc = simplify_doc(doc)
389
- for key, value in doc.items():
390
- prop = {"name": key, "type": type(value).__name__}
391
- if prop not in properties:
392
- properties.append(prop)
393
-
394
- collection_schema.append(
395
- {
488
+ properties_list = []
489
+ example_doc = None
490
+
491
+ def simplify_doc(doc: Any) -> Any:
492
+ if isinstance(doc, list) and len(doc) > 0:
493
+ return [simplify_doc(doc[0])]
494
+ if isinstance(doc, dict):
495
+ return {k: simplify_doc(v) for k, v in doc.items()}
496
+ return doc
497
+
498
+ for doc in self.db.aql.execute(sample_query): # type: ignore
499
+ if example_doc is None:
500
+ example_doc = simplify_doc(doc)
501
+ for key, value in doc.items():
502
+ prop = {"name": key, "type": type(value).__name__}
503
+ if prop not in properties_list:
504
+ properties_list.append(prop)
505
+
506
+ collection_schema.append(
507
+ {
508
+ "collection_name": col_name,
509
+ "collection_type": col_type,
510
+ f"{col_type}_properties": properties_list,
511
+ f"example_{col_type}": example_doc,
512
+ }
513
+ )
514
+ else:
515
+ # Basic info + from/to for edges only
516
+ collection_info = {
396
517
  "collection_name": col_name,
397
518
  "collection_type": col_type,
398
- f"{col_type}_properties": properties,
399
- f"example_{col_type}": example_doc,
400
519
  }
401
- )
520
+ if col_type == "edge":
521
+ # Get a sample edge to extract from/to fields
522
+ sample_edge = next(
523
+ self.db.aql.execute( # type: ignore
524
+ f"FOR e IN {col_name} LIMIT 1 RETURN e"
525
+ ),
526
+ None,
527
+ )
528
+ if sample_edge:
529
+ collection_info["from_collection"] = sample_edge[
530
+ "_from"
531
+ ].split("/")[0]
532
+ collection_info["to_collection"] = sample_edge["_to"].split(
533
+ "/"
534
+ )[0]
535
+
536
+ collection_schema.append(collection_info)
402
537
 
403
538
  schema = {
404
539
  "Graph Schema": graph_schema,
@@ -406,10 +541,41 @@ class ArangoChatAgent(ChatAgent):
406
541
  }
407
542
  schema_str = json.dumps(schema, indent=2)
408
543
  logger.warning(f"Schema retrieved:\n{schema_str}")
409
- # save schema to file "logs/arangoo-schema.json"
410
544
  with open("logs/arango-schema.json", "w") as f:
411
545
  f.write(schema_str)
412
- self.config.kg_schema = schema # type: ignore
546
+ if (n_fields := count_fields(schema)) > self.config.max_schema_fields:
547
+ logger.warning(
548
+ f"""
549
+ Schema has {n_fields} fields, which exceeds the maximum of
550
+ {self.config.max_schema_fields}. Showing a trimmed version
551
+ that only includes edge info and no other properties.
552
+ """
553
+ )
554
+ schema = trim_schema(schema)
555
+ n_fields = count_fields(schema)
556
+ logger.warning(f"Schema trimmed down to {n_fields} fields.")
557
+ schema_str = (
558
+ json.dumps(schema)
559
+ + "\n"
560
+ + f"""
561
+
562
+ CAUTION: The requested schema was too large, so
563
+ the schema has been trimmed down to show only all collection names,
564
+ their types,
565
+ and edge relationships (from/to collections) without any properties.
566
+ To find out more about the schema, you can EITHER:
567
+ - Use the `{arango_schema_tool_name}` tool again with the
568
+ `properties` arg set to True, and `collections` arg set to
569
+ specific collections you want to know more about, OR
570
+ - Use the `{aql_retrieval_tool_name}` tool to learn more about
571
+ the schema by querying the database.
572
+
573
+ """
574
+ )
575
+ if msg is None:
576
+ self.config.kg_schema = schema_str
577
+ return schema_str
578
+ self.config.kg_schema = schema
413
579
  return schema
414
580
 
415
581
  except Exception as e:
@@ -432,9 +598,10 @@ class ArangoChatAgent(ChatAgent):
432
598
 
433
599
  super().__init__(self.config)
434
600
  # Note we are enabling GraphSchemaTool regardless of whether
435
- # self.config.use_schema_tools is True or False, because
601
+ # self.config.prepopulate_schema is True or False, because
436
602
  # even when schema provided, the agent may later want to get the schema,
437
- # e.g. if the db evolves, or if it needs to bring in the schema
603
+ # e.g. if the db evolves, or schema was trimmed due to size, or
604
+ # if it needs to bring in the schema into recent context.
438
605
 
439
606
  self.enable_message(
440
607
  [
@@ -454,7 +621,7 @@ class ArangoChatAgent(ChatAgent):
454
621
  assert isinstance(self.config, ArangoChatAgentConfig)
455
622
  return (
456
623
  SCHEMA_TOOLS_SYS_MSG
457
- if self.config.use_schema_tools
624
+ if not self.config.prepopulate_schema
458
625
  else SCHEMA_PROVIDED_SYS_MSG.format(schema=self.arango_schema_tool(None))
459
626
  )
460
627
 
@@ -9,7 +9,7 @@ done_tool_name = DoneTool.default_value("request")
9
9
 
10
10
  arango_schema_tool_description = f"""
11
11
  `{arango_schema_tool_name}` tool/function-call to find the schema
12
- of the graph database, i.e. get all the collections
12
+ of the graph database, or for some SPECIFIC collections, i.e. get information on
13
13
  (document and edge), their attributes, and graph definitions available in your
14
14
  ArangoDB database. You MUST use this tool BEFORE attempting to use the
15
15
  `{aql_retrieval_tool_name}` tool/function-call, to ensure that you are using the
@@ -18,7 +18,8 @@ correct collection names and attributes in your `{aql_retrieval_tool_name}` tool
18
18
 
19
19
  aql_retrieval_tool_description = f"""
20
20
  `{aql_retrieval_tool_name}` tool/function-call to retrieve information from
21
- the database using AQL (ArangoDB Query Language) queries.
21
+ the database using AQL (ArangoDB Query Language) queries, to answer
22
+ the user's questions, OR for you to learn more about the SCHEMA of the database.
22
23
  """
23
24
 
24
25
  aql_creation_tool_description = f"""
@@ -26,6 +27,29 @@ aql_creation_tool_description = f"""
26
27
  documents/edges in the database.
27
28
  """
28
29
 
30
+ aql_retrieval_query_example = """
31
+ EXAMPLE:
32
+ Suppose you are asked this question "Does Bob have a father?".
33
+ Then you will go through the following steps, where YOU indicates
34
+ the message YOU will be sending, and RESULTS indicates the RESULTS
35
+ you will receive from the helper executing the query:
36
+
37
+ 1. YOU:
38
+ {{ "request": "aql_retrieval_tool",
39
+ "aql_query": "FOR v, e, p in ... [query truncated for brevity]..."}}
40
+
41
+ 2. RESULTS:
42
+ [.. results from the query...]
43
+ 3. YOU: [ since results were not satisfactory, you try ANOTHER query]
44
+ {{ "request": "aql_retrieval_tool",
45
+ "aql_query": "blah blah ... [query truncated for brevity]..."}}
46
+ }}
47
+ 4. RESULTS:
48
+ [.. results from the query...]
49
+ 5. YOU: [ now you have the answer, you can generate your response ]
50
+ The answer is YES, Bob has a father, and his name is John.
51
+ """
52
+
29
53
  aql_query_instructions = """
30
54
  When writing AQL queries:
31
55
  1. Use the exact property names shown in the schema
@@ -63,6 +87,7 @@ REMEMBER:
63
87
  with your response. DO NOT MAKE UP RESULTS FROM A TOOL!
64
88
  [3] YOU MUST NOT ANSWER queries from your OWN KNOWLEDGE; ALWAYS RELY ON
65
89
  the result of a TOOL/FUNCTION to compose your response.
90
+ [4] Use ONLY ONE TOOL/FUNCTION at a TIME!
66
91
  """
67
92
  # sys msg to use when schema already provided initially,
68
93
  # so agent should not use schema tool
@@ -77,6 +102,7 @@ and their attribute keys available in your ArangoDB database.
77
102
  {{schema}}
78
103
  === END SCHEMA ===
79
104
 
105
+
80
106
  To help with the user's question or database update/creation request,
81
107
  you have access to these tools:
82
108
 
@@ -84,10 +110,6 @@ you have access to these tools:
84
110
 
85
111
  - {aql_creation_tool_description}
86
112
 
87
- Since the schema has been provided, you may not need to use the tool below,
88
- but you may use it if you need to remind yourself about the schema:
89
-
90
- - {arango_schema_tool_description}
91
113
 
92
114
  {tool_result_instruction}
93
115
  """
@@ -113,27 +135,34 @@ DEFAULT_ARANGO_CHAT_SYSTEM_MESSAGE = f"""
113
135
  {{mode}}
114
136
 
115
137
  You do not need to be able to answer a question with just one query.
116
- You could make a sequence of AQL queries to find the answer to the question.
138
+ You can make a query, WAIT for the result,
139
+ THEN make ANOTHER query, WAIT for result,
140
+ THEN make ANOTHER query, and so on, until you have the answer.
117
141
 
118
142
  {aql_query_instructions}
119
143
 
120
144
  RETRY-SUGGESTIONS:
121
145
  If you receive a null or other unexpected result,
122
146
  (a) make sure you use the available TOOLs correctly,
123
- (b) USE `{arango_schema_tool_name}` tool/function-call to get all collections,
124
- their attributes and graph definitions available in your ArangoDB database.
147
+ (b) learn more about the schema using EITHER:
148
+ - `{arango_schema_tool_name}` tool/function-call to find properties of specific
149
+ collections or other parts of the schema, OR
150
+ - `{aql_retrieval_tool_name}` tool/function-call to use AQL queries to
151
+ find specific parts of the schema.
125
152
  (c) Collection names are CASE-SENSITIVE -- make sure you adhere to the exact
126
153
  collection name you found in the schema.
127
154
  (d) see if you have made an assumption in your AQL query, and try another way,
128
155
  or use `{aql_retrieval_tool_name}` to explore the database contents before
129
156
  submitting your final query.
130
- (f) Try APPROXIMATE or PARTIAL MATCHES to strings in the user's query,
157
+ (e) Try APPROXIMATE or PARTIAL MATCHES to strings in the user's query,
131
158
  e.g. user may ask about "Godfather" instead of "The Godfather",
132
159
  or try using CASE-INSENSITIVE MATCHES.
133
160
 
134
161
  Start by asking what the user needs help with.
135
162
 
136
163
  {tool_result_instruction}
164
+
165
+ {aql_retrieval_query_example}
137
166
  """
138
167
 
139
168
  ADDRESSING_INSTRUCTION = """