shotgun-sh 0.2.8.dev2__py3-none-any.whl → 0.3.3.dev1__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.
Files changed (175) hide show
  1. shotgun/agents/agent_manager.py +382 -60
  2. shotgun/agents/common.py +15 -9
  3. shotgun/agents/config/README.md +89 -0
  4. shotgun/agents/config/__init__.py +10 -1
  5. shotgun/agents/config/constants.py +0 -6
  6. shotgun/agents/config/manager.py +383 -82
  7. shotgun/agents/config/models.py +122 -18
  8. shotgun/agents/config/provider.py +81 -15
  9. shotgun/agents/config/streaming_test.py +119 -0
  10. shotgun/agents/context_analyzer/__init__.py +28 -0
  11. shotgun/agents/context_analyzer/analyzer.py +475 -0
  12. shotgun/agents/context_analyzer/constants.py +9 -0
  13. shotgun/agents/context_analyzer/formatter.py +115 -0
  14. shotgun/agents/context_analyzer/models.py +212 -0
  15. shotgun/agents/conversation/__init__.py +18 -0
  16. shotgun/agents/conversation/filters.py +164 -0
  17. shotgun/agents/conversation/history/chunking.py +278 -0
  18. shotgun/agents/{history → conversation/history}/compaction.py +36 -5
  19. shotgun/agents/{history → conversation/history}/constants.py +5 -0
  20. shotgun/agents/conversation/history/file_content_deduplication.py +216 -0
  21. shotgun/agents/{history → conversation/history}/history_processors.py +380 -8
  22. shotgun/agents/{history → conversation/history}/token_counting/anthropic.py +25 -1
  23. shotgun/agents/{history → conversation/history}/token_counting/base.py +14 -3
  24. shotgun/agents/{history → conversation/history}/token_counting/openai.py +11 -1
  25. shotgun/agents/{history → conversation/history}/token_counting/sentencepiece_counter.py +8 -0
  26. shotgun/agents/{history → conversation/history}/token_counting/tokenizer_cache.py +3 -1
  27. shotgun/agents/{history → conversation/history}/token_counting/utils.py +0 -3
  28. shotgun/agents/{conversation_manager.py → conversation/manager.py} +36 -20
  29. shotgun/agents/{conversation_history.py → conversation/models.py} +8 -92
  30. shotgun/agents/error/__init__.py +11 -0
  31. shotgun/agents/error/models.py +19 -0
  32. shotgun/agents/export.py +2 -2
  33. shotgun/agents/plan.py +2 -2
  34. shotgun/agents/research.py +3 -3
  35. shotgun/agents/runner.py +230 -0
  36. shotgun/agents/specify.py +2 -2
  37. shotgun/agents/tasks.py +2 -2
  38. shotgun/agents/tools/codebase/codebase_shell.py +6 -0
  39. shotgun/agents/tools/codebase/directory_lister.py +6 -0
  40. shotgun/agents/tools/codebase/file_read.py +11 -2
  41. shotgun/agents/tools/codebase/query_graph.py +6 -0
  42. shotgun/agents/tools/codebase/retrieve_code.py +6 -0
  43. shotgun/agents/tools/file_management.py +27 -7
  44. shotgun/agents/tools/registry.py +217 -0
  45. shotgun/agents/tools/web_search/__init__.py +8 -8
  46. shotgun/agents/tools/web_search/anthropic.py +8 -2
  47. shotgun/agents/tools/web_search/gemini.py +7 -1
  48. shotgun/agents/tools/web_search/openai.py +8 -2
  49. shotgun/agents/tools/web_search/utils.py +2 -2
  50. shotgun/agents/usage_manager.py +16 -11
  51. shotgun/api_endpoints.py +7 -3
  52. shotgun/build_constants.py +2 -2
  53. shotgun/cli/clear.py +53 -0
  54. shotgun/cli/compact.py +188 -0
  55. shotgun/cli/config.py +8 -5
  56. shotgun/cli/context.py +154 -0
  57. shotgun/cli/error_handler.py +24 -0
  58. shotgun/cli/export.py +34 -34
  59. shotgun/cli/feedback.py +4 -2
  60. shotgun/cli/models.py +1 -0
  61. shotgun/cli/plan.py +34 -34
  62. shotgun/cli/research.py +18 -10
  63. shotgun/cli/spec/__init__.py +5 -0
  64. shotgun/cli/spec/backup.py +81 -0
  65. shotgun/cli/spec/commands.py +132 -0
  66. shotgun/cli/spec/models.py +48 -0
  67. shotgun/cli/spec/pull_service.py +219 -0
  68. shotgun/cli/specify.py +20 -19
  69. shotgun/cli/tasks.py +34 -34
  70. shotgun/cli/update.py +16 -2
  71. shotgun/codebase/core/change_detector.py +5 -3
  72. shotgun/codebase/core/code_retrieval.py +4 -2
  73. shotgun/codebase/core/ingestor.py +163 -15
  74. shotgun/codebase/core/manager.py +13 -4
  75. shotgun/codebase/core/nl_query.py +1 -1
  76. shotgun/codebase/models.py +2 -0
  77. shotgun/exceptions.py +357 -0
  78. shotgun/llm_proxy/__init__.py +17 -0
  79. shotgun/llm_proxy/client.py +215 -0
  80. shotgun/llm_proxy/models.py +137 -0
  81. shotgun/logging_config.py +60 -27
  82. shotgun/main.py +77 -11
  83. shotgun/posthog_telemetry.py +38 -29
  84. shotgun/prompts/agents/partials/common_agent_system_prompt.j2 +28 -2
  85. shotgun/prompts/agents/partials/interactive_mode.j2 +3 -3
  86. shotgun/prompts/agents/plan.j2 +16 -0
  87. shotgun/prompts/agents/research.j2 +16 -3
  88. shotgun/prompts/agents/specify.j2 +54 -1
  89. shotgun/prompts/agents/state/system_state.j2 +0 -2
  90. shotgun/prompts/agents/tasks.j2 +16 -0
  91. shotgun/prompts/history/chunk_summarization.j2 +34 -0
  92. shotgun/prompts/history/combine_summaries.j2 +53 -0
  93. shotgun/sdk/codebase.py +14 -3
  94. shotgun/sentry_telemetry.py +163 -16
  95. shotgun/settings.py +243 -0
  96. shotgun/shotgun_web/__init__.py +67 -1
  97. shotgun/shotgun_web/client.py +42 -1
  98. shotgun/shotgun_web/constants.py +46 -0
  99. shotgun/shotgun_web/exceptions.py +29 -0
  100. shotgun/shotgun_web/models.py +390 -0
  101. shotgun/shotgun_web/shared_specs/__init__.py +32 -0
  102. shotgun/shotgun_web/shared_specs/file_scanner.py +175 -0
  103. shotgun/shotgun_web/shared_specs/hasher.py +83 -0
  104. shotgun/shotgun_web/shared_specs/models.py +71 -0
  105. shotgun/shotgun_web/shared_specs/upload_pipeline.py +329 -0
  106. shotgun/shotgun_web/shared_specs/utils.py +34 -0
  107. shotgun/shotgun_web/specs_client.py +703 -0
  108. shotgun/shotgun_web/supabase_client.py +31 -0
  109. shotgun/telemetry.py +10 -33
  110. shotgun/tui/app.py +310 -46
  111. shotgun/tui/commands/__init__.py +1 -1
  112. shotgun/tui/components/context_indicator.py +179 -0
  113. shotgun/tui/components/mode_indicator.py +70 -0
  114. shotgun/tui/components/status_bar.py +48 -0
  115. shotgun/tui/containers.py +91 -0
  116. shotgun/tui/dependencies.py +39 -0
  117. shotgun/tui/layout.py +5 -0
  118. shotgun/tui/protocols.py +45 -0
  119. shotgun/tui/screens/chat/__init__.py +5 -0
  120. shotgun/tui/screens/chat/chat.tcss +54 -0
  121. shotgun/tui/screens/chat/chat_screen.py +1531 -0
  122. shotgun/tui/screens/chat/codebase_index_prompt_screen.py +243 -0
  123. shotgun/tui/screens/chat/codebase_index_selection.py +12 -0
  124. shotgun/tui/screens/chat/help_text.py +40 -0
  125. shotgun/tui/screens/chat/prompt_history.py +48 -0
  126. shotgun/tui/screens/chat.tcss +11 -0
  127. shotgun/tui/screens/chat_screen/command_providers.py +91 -4
  128. shotgun/tui/screens/chat_screen/hint_message.py +76 -1
  129. shotgun/tui/screens/chat_screen/history/__init__.py +22 -0
  130. shotgun/tui/screens/chat_screen/history/agent_response.py +66 -0
  131. shotgun/tui/screens/chat_screen/history/chat_history.py +115 -0
  132. shotgun/tui/screens/chat_screen/history/formatters.py +115 -0
  133. shotgun/tui/screens/chat_screen/history/partial_response.py +43 -0
  134. shotgun/tui/screens/chat_screen/history/user_question.py +42 -0
  135. shotgun/tui/screens/confirmation_dialog.py +191 -0
  136. shotgun/tui/screens/directory_setup.py +45 -41
  137. shotgun/tui/screens/feedback.py +14 -7
  138. shotgun/tui/screens/github_issue.py +111 -0
  139. shotgun/tui/screens/model_picker.py +77 -32
  140. shotgun/tui/screens/onboarding.py +580 -0
  141. shotgun/tui/screens/pipx_migration.py +205 -0
  142. shotgun/tui/screens/provider_config.py +116 -35
  143. shotgun/tui/screens/shared_specs/__init__.py +21 -0
  144. shotgun/tui/screens/shared_specs/create_spec_dialog.py +273 -0
  145. shotgun/tui/screens/shared_specs/models.py +56 -0
  146. shotgun/tui/screens/shared_specs/share_specs_dialog.py +390 -0
  147. shotgun/tui/screens/shared_specs/upload_progress_screen.py +452 -0
  148. shotgun/tui/screens/shotgun_auth.py +112 -18
  149. shotgun/tui/screens/spec_pull.py +288 -0
  150. shotgun/tui/screens/welcome.py +137 -11
  151. shotgun/tui/services/__init__.py +5 -0
  152. shotgun/tui/services/conversation_service.py +187 -0
  153. shotgun/tui/state/__init__.py +7 -0
  154. shotgun/tui/state/processing_state.py +185 -0
  155. shotgun/tui/utils/mode_progress.py +14 -7
  156. shotgun/tui/widgets/__init__.py +5 -0
  157. shotgun/tui/widgets/widget_coordinator.py +263 -0
  158. shotgun/utils/file_system_utils.py +22 -2
  159. shotgun/utils/marketing.py +110 -0
  160. shotgun/utils/update_checker.py +69 -14
  161. shotgun_sh-0.3.3.dev1.dist-info/METADATA +472 -0
  162. shotgun_sh-0.3.3.dev1.dist-info/RECORD +229 -0
  163. {shotgun_sh-0.2.8.dev2.dist-info → shotgun_sh-0.3.3.dev1.dist-info}/WHEEL +1 -1
  164. {shotgun_sh-0.2.8.dev2.dist-info → shotgun_sh-0.3.3.dev1.dist-info}/entry_points.txt +1 -0
  165. {shotgun_sh-0.2.8.dev2.dist-info → shotgun_sh-0.3.3.dev1.dist-info}/licenses/LICENSE +1 -1
  166. shotgun/tui/screens/chat.py +0 -996
  167. shotgun/tui/screens/chat_screen/history.py +0 -335
  168. shotgun_sh-0.2.8.dev2.dist-info/METADATA +0 -126
  169. shotgun_sh-0.2.8.dev2.dist-info/RECORD +0 -155
  170. /shotgun/agents/{history → conversation/history}/__init__.py +0 -0
  171. /shotgun/agents/{history → conversation/history}/context_extraction.py +0 -0
  172. /shotgun/agents/{history → conversation/history}/history_building.py +0 -0
  173. /shotgun/agents/{history → conversation/history}/message_utils.py +0 -0
  174. /shotgun/agents/{history → conversation/history}/token_counting/__init__.py +0 -0
  175. /shotgun/agents/{history → conversation/history}/token_estimation.py +0 -0
@@ -2,9 +2,11 @@
2
2
 
3
3
  from pathlib import Path
4
4
 
5
+ import aiofiles
5
6
  from pydantic_ai import RunContext
6
7
 
7
8
  from shotgun.agents.models import AgentDeps
9
+ from shotgun.agents.tools.registry import ToolCategory, register_tool
8
10
  from shotgun.codebase.core.language_config import get_language_config
9
11
  from shotgun.logging_config import get_logger
10
12
 
@@ -13,6 +15,11 @@ from .models import FileReadResult
13
15
  logger = get_logger(__name__)
14
16
 
15
17
 
18
+ @register_tool(
19
+ category=ToolCategory.CODEBASE_UNDERSTANDING,
20
+ display_text="Reading file",
21
+ key_arg="file_path",
22
+ )
16
23
  async def file_read(
17
24
  ctx: RunContext[AgentDeps], graph_id: str, file_path: str
18
25
  ) -> FileReadResult:
@@ -87,7 +94,8 @@ async def file_read(
87
94
  # Read file contents
88
95
  encoding_used = "utf-8"
89
96
  try:
90
- content = full_file_path.read_text(encoding="utf-8")
97
+ async with aiofiles.open(full_file_path, encoding="utf-8") as f:
98
+ content = await f.read()
91
99
  size_bytes = full_file_path.stat().st_size
92
100
 
93
101
  logger.debug(
@@ -113,7 +121,8 @@ async def file_read(
113
121
  try:
114
122
  # Try with different encoding
115
123
  encoding_used = "latin-1"
116
- content = full_file_path.read_text(encoding="latin-1")
124
+ async with aiofiles.open(full_file_path, encoding="latin-1") as f:
125
+ content = await f.read()
117
126
  size_bytes = full_file_path.stat().st_size
118
127
 
119
128
  # Detect language from file extension
@@ -3,6 +3,7 @@
3
3
  from pydantic_ai import RunContext
4
4
 
5
5
  from shotgun.agents.models import AgentDeps
6
+ from shotgun.agents.tools.registry import ToolCategory, register_tool
6
7
  from shotgun.codebase.models import QueryType
7
8
  from shotgun.logging_config import get_logger
8
9
 
@@ -11,6 +12,11 @@ from .models import QueryGraphResult
11
12
  logger = get_logger(__name__)
12
13
 
13
14
 
15
+ @register_tool(
16
+ category=ToolCategory.CODEBASE_UNDERSTANDING,
17
+ display_text="Querying code",
18
+ key_arg="query",
19
+ )
14
20
  async def query_graph(
15
21
  ctx: RunContext[AgentDeps], graph_id: str, query: str
16
22
  ) -> QueryGraphResult:
@@ -5,6 +5,7 @@ from pathlib import Path
5
5
  from pydantic_ai import RunContext
6
6
 
7
7
  from shotgun.agents.models import AgentDeps
8
+ from shotgun.agents.tools.registry import ToolCategory, register_tool
8
9
  from shotgun.codebase.core.code_retrieval import retrieve_code_by_qualified_name
9
10
  from shotgun.codebase.core.language_config import get_language_config
10
11
  from shotgun.logging_config import get_logger
@@ -14,6 +15,11 @@ from .models import CodeSnippetResult
14
15
  logger = get_logger(__name__)
15
16
 
16
17
 
18
+ @register_tool(
19
+ category=ToolCategory.CODEBASE_UNDERSTANDING,
20
+ display_text="Retrieving code",
21
+ key_arg="qualified_name",
22
+ )
17
23
  async def retrieve_code(
18
24
  ctx: RunContext[AgentDeps], graph_id: str, qualified_name: str
19
25
  ) -> CodeSnippetResult:
@@ -6,9 +6,12 @@ These tools are restricted to the .shotgun directory for security.
6
6
  from pathlib import Path
7
7
  from typing import Literal
8
8
 
9
+ import aiofiles
10
+ import aiofiles.os
9
11
  from pydantic_ai import RunContext
10
12
 
11
13
  from shotgun.agents.models import AgentDeps, AgentType, FileOperationType
14
+ from shotgun.agents.tools.registry import ToolCategory, register_tool
12
15
  from shotgun.logging_config import get_logger
13
16
  from shotgun.utils.file_system_utils import get_shotgun_base_path
14
17
 
@@ -157,6 +160,11 @@ def _validate_shotgun_path(filename: str) -> Path:
157
160
  return full_path
158
161
 
159
162
 
163
+ @register_tool(
164
+ category=ToolCategory.ARTIFACT_MANAGEMENT,
165
+ display_text="Reading file",
166
+ key_arg="filename",
167
+ )
160
168
  async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
161
169
  """Read a file from the .shotgun directory.
162
170
 
@@ -175,10 +183,11 @@ async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
175
183
  try:
176
184
  file_path = _validate_shotgun_path(filename)
177
185
 
178
- if not file_path.exists():
186
+ if not await aiofiles.os.path.exists(file_path):
179
187
  raise FileNotFoundError(f"File not found: {filename}")
180
188
 
181
- content = file_path.read_text(encoding="utf-8")
189
+ async with aiofiles.open(file_path, encoding="utf-8") as f:
190
+ content = await f.read()
182
191
  logger.debug("📄 Read %d characters from %s", len(content), filename)
183
192
  return content
184
193
 
@@ -188,6 +197,11 @@ async def read_file(ctx: RunContext[AgentDeps], filename: str) -> str:
188
197
  return error_msg
189
198
 
190
199
 
200
+ @register_tool(
201
+ category=ToolCategory.ARTIFACT_MANAGEMENT,
202
+ display_text="Writing file",
203
+ key_arg="filename",
204
+ )
191
205
  async def write_file(
192
206
  ctx: RunContext[AgentDeps],
193
207
  filename: str,
@@ -222,21 +236,22 @@ async def write_file(
222
236
  else:
223
237
  operation = (
224
238
  FileOperationType.CREATED
225
- if not file_path.exists()
239
+ if not await aiofiles.os.path.exists(file_path)
226
240
  else FileOperationType.UPDATED
227
241
  )
228
242
 
229
243
  # Ensure parent directory exists
230
- file_path.parent.mkdir(parents=True, exist_ok=True)
244
+ await aiofiles.os.makedirs(file_path.parent, exist_ok=True)
231
245
 
232
246
  # Write content
233
247
  if mode == "a":
234
- with open(file_path, "a", encoding="utf-8") as f:
235
- f.write(content)
248
+ async with aiofiles.open(file_path, "a", encoding="utf-8") as f:
249
+ await f.write(content)
236
250
  logger.debug("📄 Appended %d characters to %s", len(content), filename)
237
251
  result = f"Successfully appended {len(content)} characters to {filename}"
238
252
  else:
239
- file_path.write_text(content, encoding="utf-8")
253
+ async with aiofiles.open(file_path, "w", encoding="utf-8") as f:
254
+ await f.write(content)
240
255
  logger.debug("📄 Wrote %d characters to %s", len(content), filename)
241
256
  result = f"Successfully wrote {len(content)} characters to {filename}"
242
257
 
@@ -251,6 +266,11 @@ async def write_file(
251
266
  return error_msg
252
267
 
253
268
 
269
+ @register_tool(
270
+ category=ToolCategory.ARTIFACT_MANAGEMENT,
271
+ display_text="Appending to file",
272
+ key_arg="filename",
273
+ )
254
274
  async def append_file(ctx: RunContext[AgentDeps], filename: str, content: str) -> str:
255
275
  """Append content to a file in the .shotgun directory.
256
276
 
@@ -0,0 +1,217 @@
1
+ """Tool category registry using decorators for automatic registration.
2
+
3
+ This module provides a decorator-based system for categorizing tools used by agents.
4
+ Tools can be decorated with @register_tool to automatically register their category,
5
+ which is then used by the context analyzer to break down token usage by tool type.
6
+
7
+ It also provides a display registry system for tool formatting in the TUI, allowing
8
+ tools to declare how they should be displayed when streaming.
9
+ """
10
+
11
+ from collections.abc import Callable
12
+ from enum import StrEnum
13
+ from typing import TypeVar, overload
14
+
15
+ import sentry_sdk
16
+ from pydantic import BaseModel
17
+
18
+ from shotgun.logging_config import get_logger
19
+
20
+ logger = get_logger(__name__)
21
+
22
+ # Type variable for decorated functions
23
+ F = TypeVar("F", bound=Callable[..., object])
24
+
25
+
26
+ class ToolCategory(StrEnum):
27
+ """Categories for agent tools used in context analysis."""
28
+
29
+ CODEBASE_UNDERSTANDING = "codebase_understanding"
30
+ ARTIFACT_MANAGEMENT = "artifact_management"
31
+ WEB_RESEARCH = "web_research"
32
+ AGENT_RESPONSE = "agent_response"
33
+ UNKNOWN = "unknown"
34
+
35
+
36
+ class ToolDisplayConfig(BaseModel):
37
+ """Configuration for how a tool should be displayed in the TUI.
38
+
39
+ Attributes:
40
+ display_text: Text to show (e.g., "Reading file", "Querying code")
41
+ key_arg: Primary argument to extract from tool args for display
42
+ hide: Whether to completely hide this tool call from the UI
43
+ """
44
+
45
+ display_text: str
46
+ key_arg: str
47
+ hide: bool = False
48
+
49
+
50
+ # Global registry mapping tool names to categories
51
+ _TOOL_REGISTRY: dict[str, ToolCategory] = {}
52
+
53
+ # Global registry mapping tool names to display configs
54
+ _TOOL_DISPLAY_REGISTRY: dict[str, ToolDisplayConfig] = {}
55
+
56
+
57
+ @overload
58
+ def register_tool(
59
+ category: ToolCategory,
60
+ display_text: str,
61
+ key_arg: str,
62
+ ) -> Callable[[F], F]: ...
63
+
64
+
65
+ @overload
66
+ def register_tool(
67
+ category: ToolCategory,
68
+ display_text: str,
69
+ key_arg: str,
70
+ *,
71
+ hide: bool,
72
+ ) -> Callable[[F], F]: ...
73
+
74
+
75
+ def register_tool(
76
+ category: ToolCategory,
77
+ display_text: str,
78
+ key_arg: str,
79
+ *,
80
+ hide: bool = False,
81
+ ) -> Callable[[F], F]:
82
+ """Decorator to register a tool's category and display configuration.
83
+
84
+ Args:
85
+ category: The ToolCategory enum value for this tool
86
+ display_text: Text to show (e.g., "Reading file", "Querying code")
87
+ key_arg: Primary argument name to extract for display (e.g., "query", "filename")
88
+ hide: Whether to hide this tool call completely from the UI (default: False)
89
+
90
+ Returns:
91
+ Decorator function that registers the tool and returns it unchanged
92
+
93
+ Display Format:
94
+ - When key_arg value is missing: Shows just display_text (e.g., "Reading file")
95
+ - When key_arg value is present: Shows "display_text: key_arg_value" (e.g., "Reading file: foo.py")
96
+
97
+ Example:
98
+ @register_tool(
99
+ category=ToolCategory.CODEBASE_UNDERSTANDING,
100
+ display_text="Querying code",
101
+ key_arg="query",
102
+ )
103
+ async def query_graph(ctx: RunContext[AgentDeps], query: str) -> str:
104
+ ...
105
+ """
106
+
107
+ def decorator(func: F) -> F:
108
+ tool_name = func.__name__
109
+ _TOOL_REGISTRY[tool_name] = category
110
+ logger.debug(f"Registered tool '{tool_name}' as category '{category.value}'")
111
+
112
+ # Register display config
113
+ config = ToolDisplayConfig(
114
+ display_text=display_text,
115
+ key_arg=key_arg,
116
+ hide=hide,
117
+ )
118
+ _TOOL_DISPLAY_REGISTRY[tool_name] = config
119
+ logger.debug(f"Registered display config for tool '{tool_name}'")
120
+
121
+ return func
122
+
123
+ return decorator
124
+
125
+
126
+ # Backwards compatibility alias
127
+ tool_category = register_tool
128
+
129
+
130
+ def get_tool_category(tool_name: str) -> ToolCategory:
131
+ """Get category for a tool, logging unknown tools to Sentry.
132
+
133
+ Args:
134
+ tool_name: Name of the tool to look up
135
+
136
+ Returns:
137
+ ToolCategory enum value for the tool, or UNKNOWN if not registered
138
+ """
139
+ category = _TOOL_REGISTRY.get(tool_name)
140
+
141
+ if category is None:
142
+ logger.warning(f"Unknown tool encountered in context analysis: {tool_name}")
143
+ sentry_sdk.capture_message(
144
+ f"Unknown tool in context analysis: {tool_name}",
145
+ level="warning",
146
+ extras={"tool_name": tool_name},
147
+ )
148
+ return ToolCategory.UNKNOWN
149
+
150
+ return category
151
+
152
+
153
+ def register_special_tool(tool_name: str, category: ToolCategory) -> None:
154
+ """Register a special tool that doesn't have a decorator.
155
+
156
+ Used for tools like 'final_result' that aren't actual Python functions
157
+ but need to be categorized.
158
+
159
+ Args:
160
+ tool_name: Name of the special tool
161
+ category: Category to assign to this tool
162
+ """
163
+ _TOOL_REGISTRY[tool_name] = category
164
+ logger.debug(
165
+ f"Registered special tool '{tool_name}' as category '{category.value}'"
166
+ )
167
+
168
+
169
+ def get_tool_display_config(tool_name: str) -> ToolDisplayConfig | None:
170
+ """Get display configuration for a tool.
171
+
172
+ Args:
173
+ tool_name: Name of the tool to look up
174
+
175
+ Returns:
176
+ ToolDisplayConfig for the tool, or None if not registered
177
+ """
178
+ return _TOOL_DISPLAY_REGISTRY.get(tool_name)
179
+
180
+
181
+ def register_tool_display(
182
+ tool_name: str,
183
+ display_text: str,
184
+ key_arg: str,
185
+ *,
186
+ hide: bool = False,
187
+ ) -> None:
188
+ """Register a display config for a special tool that doesn't have a decorator.
189
+
190
+ Used for tools like 'final_result' or builtin tools that aren't actual Python functions.
191
+
192
+ Args:
193
+ tool_name: Name of the special tool
194
+ display_text: Text to show (e.g., "Reading file", "Querying code")
195
+ key_arg: Primary argument name to extract for display
196
+ hide: Whether to hide this tool call completely
197
+ """
198
+ config = ToolDisplayConfig(
199
+ display_text=display_text,
200
+ key_arg=key_arg,
201
+ hide=hide,
202
+ )
203
+ _TOOL_DISPLAY_REGISTRY[tool_name] = config
204
+ logger.debug(f"Registered display config for special tool '{tool_name}'")
205
+
206
+
207
+ # Register special tools that don't have decorators
208
+ register_special_tool("final_result", ToolCategory.AGENT_RESPONSE)
209
+ register_tool_display("final_result", display_text="", key_arg="", hide=True)
210
+
211
+ # Register builtin tools (tools that come from Pydantic AI or model providers)
212
+ # These don't have Python function definitions but need display formatting
213
+ register_tool_display(
214
+ "web_search",
215
+ display_text="Searching",
216
+ key_arg="query",
217
+ )
@@ -26,7 +26,7 @@ logger = get_logger(__name__)
26
26
  WebSearchTool = Callable[[str], Awaitable[str]]
27
27
 
28
28
 
29
- def get_available_web_search_tools() -> list[WebSearchTool]:
29
+ async def get_available_web_search_tools() -> list[WebSearchTool]:
30
30
  """Get list of available web search tools based on configured API keys.
31
31
 
32
32
  Works with both Shotgun Account (via LiteLLM proxy) and BYOK (individual provider keys).
@@ -43,25 +43,25 @@ def get_available_web_search_tools() -> list[WebSearchTool]:
43
43
 
44
44
  # Check if using Shotgun Account
45
45
  config_manager = get_config_manager()
46
- config = config_manager.load()
46
+ config = await config_manager.load()
47
47
  has_shotgun_key = config.shotgun.api_key is not None
48
48
 
49
49
  if has_shotgun_key:
50
50
  logger.debug("🔑 Shotgun Account - only Gemini web search available")
51
51
 
52
52
  # Gemini: Only search tool available for Shotgun Account
53
- if is_provider_available(ProviderType.GOOGLE):
53
+ if await is_provider_available(ProviderType.GOOGLE):
54
54
  logger.debug("✅ Gemini web search tool available")
55
55
  tools.append(gemini_web_search_tool)
56
56
 
57
57
  # Anthropic: Not available for Shotgun Account (Gemini-only for Shotgun)
58
- if is_provider_available(ProviderType.ANTHROPIC):
58
+ if await is_provider_available(ProviderType.ANTHROPIC):
59
59
  logger.debug(
60
60
  "⚠️ Anthropic web search requires BYOK (Shotgun Account uses Gemini only)"
61
61
  )
62
62
 
63
63
  # OpenAI: Not available for Shotgun Account (Responses API incompatible with proxy)
64
- if is_provider_available(ProviderType.OPENAI):
64
+ if await is_provider_available(ProviderType.OPENAI):
65
65
  logger.debug(
66
66
  "⚠️ OpenAI web search requires BYOK (Responses API not supported via proxy)"
67
67
  )
@@ -69,15 +69,15 @@ def get_available_web_search_tools() -> list[WebSearchTool]:
69
69
  # BYOK mode: Load all available tools based on individual provider keys
70
70
  logger.debug("🔑 BYOK mode - checking all provider web search tools")
71
71
 
72
- if is_provider_available(ProviderType.OPENAI):
72
+ if await is_provider_available(ProviderType.OPENAI):
73
73
  logger.debug("✅ OpenAI web search tool available")
74
74
  tools.append(openai_web_search_tool)
75
75
 
76
- if is_provider_available(ProviderType.ANTHROPIC):
76
+ if await is_provider_available(ProviderType.ANTHROPIC):
77
77
  logger.debug("✅ Anthropic web search tool available")
78
78
  tools.append(anthropic_web_search_tool)
79
79
 
80
- if is_provider_available(ProviderType.GOOGLE):
80
+ if await is_provider_available(ProviderType.GOOGLE):
81
81
  logger.debug("✅ Gemini web search tool available")
82
82
  tools.append(gemini_web_search_tool)
83
83
 
@@ -8,6 +8,7 @@ from shotgun.agents.config import get_provider_model
8
8
  from shotgun.agents.config.constants import MEDIUM_TEXT_8K_TOKENS
9
9
  from shotgun.agents.config.models import ProviderType
10
10
  from shotgun.agents.llm import shotgun_model_request
11
+ from shotgun.agents.tools.registry import ToolCategory, register_tool
11
12
  from shotgun.logging_config import get_logger
12
13
  from shotgun.prompts import PromptLoader
13
14
  from shotgun.utils.datetime_utils import get_datetime_context
@@ -18,6 +19,11 @@ logger = get_logger(__name__)
18
19
  prompt_loader = PromptLoader()
19
20
 
20
21
 
22
+ @register_tool(
23
+ category=ToolCategory.WEB_RESEARCH,
24
+ display_text="Searching web",
25
+ key_arg="query",
26
+ )
21
27
  async def anthropic_web_search_tool(query: str) -> str:
22
28
  """Perform a web search using Anthropic's Claude API.
23
29
 
@@ -40,7 +46,7 @@ async def anthropic_web_search_tool(query: str) -> str:
40
46
 
41
47
  # Get model configuration (supports both Shotgun and BYOK)
42
48
  try:
43
- model_config = get_provider_model(ProviderType.ANTHROPIC)
49
+ model_config = await get_provider_model(ProviderType.ANTHROPIC)
44
50
  except ValueError as e:
45
51
  error_msg = f"Anthropic API key not configured: {str(e)}"
46
52
  logger.error("❌ %s", error_msg)
@@ -135,7 +141,7 @@ async def main() -> None:
135
141
  # Check if API key is available
136
142
  try:
137
143
  if callable(get_provider_model):
138
- model_config = get_provider_model(ProviderType.ANTHROPIC)
144
+ model_config = await get_provider_model(ProviderType.ANTHROPIC)
139
145
  if not model_config.api_key:
140
146
  raise ValueError("No API key configured")
141
147
  except (ValueError, Exception):
@@ -8,6 +8,7 @@ from shotgun.agents.config import get_provider_model
8
8
  from shotgun.agents.config.constants import MEDIUM_TEXT_8K_TOKENS
9
9
  from shotgun.agents.config.models import ModelName
10
10
  from shotgun.agents.llm import shotgun_model_request
11
+ from shotgun.agents.tools.registry import ToolCategory, register_tool
11
12
  from shotgun.logging_config import get_logger
12
13
  from shotgun.prompts import PromptLoader
13
14
  from shotgun.utils.datetime_utils import get_datetime_context
@@ -18,6 +19,11 @@ logger = get_logger(__name__)
18
19
  prompt_loader = PromptLoader()
19
20
 
20
21
 
22
+ @register_tool(
23
+ category=ToolCategory.WEB_RESEARCH,
24
+ display_text="Searching web",
25
+ key_arg="query",
26
+ )
21
27
  async def gemini_web_search_tool(query: str) -> str:
22
28
  """Perform a web search using Google's Gemini API with grounding.
23
29
 
@@ -40,7 +46,7 @@ async def gemini_web_search_tool(query: str) -> str:
40
46
 
41
47
  # Get model configuration (supports both Shotgun and BYOK)
42
48
  try:
43
- model_config = get_provider_model(ModelName.GEMINI_2_5_FLASH)
49
+ model_config = await get_provider_model(ModelName.GEMINI_2_5_FLASH)
44
50
  except ValueError as e:
45
51
  error_msg = f"Gemini API key not configured: {str(e)}"
46
52
  logger.error("❌ %s", error_msg)
@@ -5,6 +5,7 @@ from opentelemetry import trace
5
5
 
6
6
  from shotgun.agents.config import get_provider_model
7
7
  from shotgun.agents.config.models import ProviderType
8
+ from shotgun.agents.tools.registry import ToolCategory, register_tool
8
9
  from shotgun.logging_config import get_logger
9
10
  from shotgun.prompts import PromptLoader
10
11
  from shotgun.utils.datetime_utils import get_datetime_context
@@ -15,6 +16,11 @@ logger = get_logger(__name__)
15
16
  prompt_loader = PromptLoader()
16
17
 
17
18
 
19
+ @register_tool(
20
+ category=ToolCategory.WEB_RESEARCH,
21
+ display_text="Searching web",
22
+ key_arg="query",
23
+ )
18
24
  async def openai_web_search_tool(query: str) -> str:
19
25
  """Perform a web search and return results.
20
26
 
@@ -37,7 +43,7 @@ async def openai_web_search_tool(query: str) -> str:
37
43
 
38
44
  # Get API key from centralized configuration
39
45
  try:
40
- model_config = get_provider_model(ProviderType.OPENAI)
46
+ model_config = await get_provider_model(ProviderType.OPENAI)
41
47
  api_key = model_config.api_key
42
48
  except ValueError as e:
43
49
  error_msg = f"OpenAI API key not configured: {str(e)}"
@@ -58,7 +64,7 @@ async def openai_web_search_tool(query: str) -> str:
58
64
  )
59
65
 
60
66
  client = AsyncOpenAI(api_key=api_key)
61
- response = await client.responses.create( # type: ignore[call-overload]
67
+ response = await client.responses.create(
62
68
  model="gpt-5-mini",
63
69
  input=[
64
70
  {"role": "user", "content": [{"type": "input_text", "text": prompt}]}
@@ -4,7 +4,7 @@ from shotgun.agents.config import get_provider_model
4
4
  from shotgun.agents.config.models import ProviderType
5
5
 
6
6
 
7
- def is_provider_available(provider: ProviderType) -> bool:
7
+ async def is_provider_available(provider: ProviderType) -> bool:
8
8
  """Check if a provider has API key configured.
9
9
 
10
10
  Args:
@@ -14,7 +14,7 @@ def is_provider_available(provider: ProviderType) -> bool:
14
14
  True if the provider has valid credentials configured (from config or env)
15
15
  """
16
16
  try:
17
- get_provider_model(provider)
17
+ await get_provider_model(provider)
18
18
  return True
19
19
  except ValueError:
20
20
  return False
@@ -6,6 +6,8 @@ from logging import getLogger
6
6
  from pathlib import Path
7
7
  from typing import TypeAlias
8
8
 
9
+ import aiofiles
10
+ import aiofiles.os
9
11
  from genai_prices import calc_price
10
12
  from pydantic import BaseModel, Field
11
13
  from pydantic_ai import RunUsage
@@ -48,9 +50,10 @@ class SessionUsageManager:
48
50
  self._model_providers: dict[ModelName, ProviderType] = {}
49
51
  self._usage_log: list[UsageLogEntry] = []
50
52
  self._usage_path: Path = get_shotgun_home() / "usage.json"
51
- self.restore_usage_state()
53
+ # Note: restore_usage_state needs to be called asynchronously after init
54
+ # Caller should use: manager = SessionUsageManager(); await manager.restore_usage_state()
52
55
 
53
- def add_usage(
56
+ async def add_usage(
54
57
  self, usage: RunUsage, *, model_name: ModelName, provider: ProviderType
55
58
  ) -> None:
56
59
  self.usage[model_name] += usage
@@ -58,7 +61,7 @@ class SessionUsageManager:
58
61
  self._usage_log.append(
59
62
  UsageLogEntry(model_name=model_name, usage=usage, provider=provider)
60
63
  )
61
- self.persist_usage_state()
64
+ await self.persist_usage_state()
62
65
 
63
66
  def get_usage_report(self) -> dict[ModelName, RunUsage]:
64
67
  return self.usage.copy()
@@ -78,7 +81,7 @@ class SessionUsageManager:
78
81
  def build_usage_hint(self) -> str | None:
79
82
  return format_usage_hint(self.get_usage_breakdown())
80
83
 
81
- def persist_usage_state(self) -> None:
84
+ async def persist_usage_state(self) -> None:
82
85
  state = UsageState(
83
86
  usage=dict(self.usage.items()),
84
87
  model_providers=self._model_providers.copy(),
@@ -86,23 +89,25 @@ class SessionUsageManager:
86
89
  )
87
90
 
88
91
  try:
89
- self._usage_path.parent.mkdir(parents=True, exist_ok=True)
90
- with self._usage_path.open("w", encoding="utf-8") as f:
91
- json.dump(state.model_dump(mode="json"), f, indent=2)
92
+ await aiofiles.os.makedirs(self._usage_path.parent, exist_ok=True)
93
+ json_content = json.dumps(state.model_dump(mode="json"), indent=2)
94
+ async with aiofiles.open(self._usage_path, "w", encoding="utf-8") as f:
95
+ await f.write(json_content)
92
96
  logger.debug("Usage state persisted to %s", self._usage_path)
93
97
  except Exception as exc:
94
98
  logger.error(
95
99
  "Failed to persist usage state to %s: %s", self._usage_path, exc
96
100
  )
97
101
 
98
- def restore_usage_state(self) -> None:
99
- if not self._usage_path.exists():
102
+ async def restore_usage_state(self) -> None:
103
+ if not await aiofiles.os.path.exists(self._usage_path):
100
104
  logger.debug("No usage state file found at %s", self._usage_path)
101
105
  return
102
106
 
103
107
  try:
104
- with self._usage_path.open(encoding="utf-8") as f:
105
- data = json.load(f)
108
+ async with aiofiles.open(self._usage_path, encoding="utf-8") as f:
109
+ content = await f.read()
110
+ data = json.loads(content)
106
111
 
107
112
  state = UsageState.model_validate(data)
108
113
  except Exception as exc:
shotgun/api_endpoints.py CHANGED
@@ -1,10 +1,14 @@
1
1
  """Shotgun backend service API endpoints and URLs."""
2
2
 
3
+ from shotgun.settings import settings
4
+
3
5
  # Shotgun Web API base URL (for authentication/subscription)
4
- # Can be overridden with environment variable
5
- SHOTGUN_WEB_BASE_URL = "https://api-219702594231.us-east4.run.app"
6
+ # Can be overridden with SHOTGUN_WEB_BASE_URL environment variable
7
+ SHOTGUN_WEB_BASE_URL = settings.api.web_base_url
8
+
6
9
  # Shotgun's LiteLLM proxy base URL (for AI model requests)
7
- LITELLM_PROXY_BASE_URL = "https://litellm-219702594231.us-east4.run.app"
10
+ # Can be overridden with SHOTGUN_ACCOUNT_LLM_BASE_URL environment variable
11
+ LITELLM_PROXY_BASE_URL = settings.api.account_llm_base_url
8
12
 
9
13
  # Provider-specific LiteLLM proxy endpoints
10
14
  LITELLM_PROXY_ANTHROPIC_BASE = f"{LITELLM_PROXY_BASE_URL}/anthropic"
@@ -8,12 +8,12 @@ DO NOT EDIT MANUALLY.
8
8
  SENTRY_DSN = 'https://2818a6d165c64eccc94cfd51ce05d6aa@o4506813296738304.ingest.us.sentry.io/4510045952409600'
9
9
 
10
10
  # PostHog configuration embedded at build time (empty strings if not provided)
11
- POSTHOG_API_KEY = ''
11
+ POSTHOG_API_KEY = 'phc_KKnChzZUKeNqZDOTJ6soCBWNQSx3vjiULdwTR9H5Mcr'
12
12
  POSTHOG_PROJECT_ID = '191396'
13
13
 
14
14
  # Logfire configuration embedded at build time (only for dev builds)
15
15
  LOGFIRE_ENABLED = 'true'
16
- LOGFIRE_TOKEN = 'pylf_v1_us_KZ5NM1pP3NwgJkbBJt6Ftdzk8mMhmrXcGJHQQgDJ1LfK'
16
+ LOGFIRE_TOKEN = 'pylf_v1_us_RwZMlJm1tX6j0PL5RWWbmZpzK2hLBNtFWStNKlySfjh8'
17
17
 
18
18
  # Build metadata
19
19
  BUILD_TIME_ENV = "production" if SENTRY_DSN else "development"