codebase-retrieval-context-engine 2.0.0__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.
- codebase_retrieval_context_engine-2.0.0.dist-info/METADATA +505 -0
- codebase_retrieval_context_engine-2.0.0.dist-info/RECORD +46 -0
- codebase_retrieval_context_engine-2.0.0.dist-info/WHEEL +4 -0
- codebase_retrieval_context_engine-2.0.0.dist-info/entry_points.txt +3 -0
- codebase_retrieval_context_engine-2.0.0.dist-info/licenses/LICENSE +201 -0
- corbell/__init__.py +6 -0
- corbell/cli/__init__.py +1 -0
- corbell/cli/commands/__init__.py +1 -0
- corbell/cli/commands/index.py +86 -0
- corbell/cli/commands/query.py +71 -0
- corbell/cli/main.py +57 -0
- corbell/core/__init__.py +1 -0
- corbell/core/constants.py +52 -0
- corbell/core/embeddings/__init__.py +6 -0
- corbell/core/embeddings/base.py +68 -0
- corbell/core/embeddings/extractor.py +201 -0
- corbell/core/embeddings/factory.py +48 -0
- corbell/core/embeddings/model.py +401 -0
- corbell/core/embeddings/search_cache.py +95 -0
- corbell/core/embeddings/sqlite_store.py +271 -0
- corbell/core/gitignore.py +76 -0
- corbell/core/graph/__init__.py +1 -0
- corbell/core/graph/builder.py +696 -0
- corbell/core/graph/method_graph.py +1077 -0
- corbell/core/graph/providers/__init__.py +6 -0
- corbell/core/graph/providers/aws_patterns.py +62 -0
- corbell/core/graph/providers/azure_patterns.py +64 -0
- corbell/core/graph/providers/gcp_patterns.py +59 -0
- corbell/core/graph/schema.py +175 -0
- corbell/core/graph/sqlite_store.py +500 -0
- corbell/core/indexing/__init__.py +1 -0
- corbell/core/indexing/builder.py +608 -0
- corbell/core/indexing/lock.py +150 -0
- corbell/core/indexing/tracker.py +245 -0
- corbell/core/llm_client.py +677 -0
- corbell/core/mcp/__init__.py +1 -0
- corbell/core/mcp/server.py +214 -0
- corbell/core/query/__init__.py +1 -0
- corbell/core/query/diagnostics.py +38 -0
- corbell/core/query/engine.py +321 -0
- corbell/core/query/enhancer.py +102 -0
- corbell/core/query/formatter.py +98 -0
- corbell/core/query/graph_expander.py +284 -0
- corbell/core/query/merger.py +171 -0
- corbell/core/query/reranker.py +131 -0
- corbell/core/workspace.py +408 -0
|
@@ -0,0 +1,677 @@
|
|
|
1
|
+
"""LLM client for Corbell — multi-provider with cloud support.
|
|
2
|
+
|
|
3
|
+
Supports local providers:
|
|
4
|
+
- ``anthropic`` — requires ``anthropic>=0.25`` and ANTHROPIC_API_KEY
|
|
5
|
+
- ``openai`` — requires ``openai>=1.0`` and OPENAI_API_KEY
|
|
6
|
+
- ``ollama`` — requires a running Ollama server (http://localhost:11434)
|
|
7
|
+
- ``google`` — requires ``google-genai>=2.7.0`` and GOOGLE_API_KEY
|
|
8
|
+
|
|
9
|
+
And cloud-hosted providers (for enterprise teams with existing cloud commitments):
|
|
10
|
+
- ``aws`` — Anthropic Claude via AWS Bedrock (boto3 + AWS credentials)
|
|
11
|
+
- ``azure`` — OpenAI GPT-4 via Azure OpenAI Service (openai + Azure endpoint)
|
|
12
|
+
- ``gcp`` — Anthropic Claude via GCP Vertex AI (google-cloud-aiplatform)
|
|
13
|
+
|
|
14
|
+
Token usage is automatically tracked in the provided TokenUsageTracker.
|
|
15
|
+
"""
|
|
16
|
+
|
|
17
|
+
from __future__ import annotations
|
|
18
|
+
|
|
19
|
+
import json
|
|
20
|
+
import os
|
|
21
|
+
import random
|
|
22
|
+
from typing import Optional, TYPE_CHECKING
|
|
23
|
+
|
|
24
|
+
if TYPE_CHECKING:
|
|
25
|
+
from corbell.core.token_tracker import TokenUsageTracker
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
class LLMClient:
|
|
29
|
+
"""Provider-agnostic LLM client for Corbell.
|
|
30
|
+
|
|
31
|
+
**Local providers** (public API keys):
|
|
32
|
+
|
|
33
|
+
.. code-block:: yaml
|
|
34
|
+
|
|
35
|
+
llm:
|
|
36
|
+
provider: anthropic
|
|
37
|
+
model: claude-sonnet-4-5-20250929
|
|
38
|
+
api_key: ${ANTHROPIC_API_KEY}
|
|
39
|
+
|
|
40
|
+
Google AI (Gemini):
|
|
41
|
+
|
|
42
|
+
.. code-block:: yaml
|
|
43
|
+
|
|
44
|
+
llm:
|
|
45
|
+
provider: google
|
|
46
|
+
model: gemini-2.5-flash
|
|
47
|
+
api_key: ${GOOGLE_API_KEY}
|
|
48
|
+
|
|
49
|
+
**Cloud providers** (enterprise API keys from your cloud console):
|
|
50
|
+
|
|
51
|
+
AWS Bedrock:
|
|
52
|
+
|
|
53
|
+
.. code-block:: yaml
|
|
54
|
+
|
|
55
|
+
llm:
|
|
56
|
+
provider: aws
|
|
57
|
+
model: anthropic.claude-sonnet-4-5-20250929-v2:0
|
|
58
|
+
aws_region: us-east-1
|
|
59
|
+
# Credentials from env: AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY
|
|
60
|
+
# or from ~/.aws/credentials profile
|
|
61
|
+
|
|
62
|
+
Azure OpenAI:
|
|
63
|
+
|
|
64
|
+
.. code-block:: yaml
|
|
65
|
+
|
|
66
|
+
llm:
|
|
67
|
+
provider: azure
|
|
68
|
+
model: gpt-4o
|
|
69
|
+
azure_endpoint: https://my-resource.openai.azure.com/
|
|
70
|
+
azure_deployment: my-gpt4o-deployment
|
|
71
|
+
azure_api_version: "2024-02-01"
|
|
72
|
+
api_key: ${AZURE_OPENAI_API_KEY}
|
|
73
|
+
|
|
74
|
+
GCP Vertex AI:
|
|
75
|
+
|
|
76
|
+
.. code-block:: yaml
|
|
77
|
+
|
|
78
|
+
llm:
|
|
79
|
+
provider: gcp
|
|
80
|
+
model: claude-3-5-sonnet@20241022
|
|
81
|
+
gcp_project: my-gcp-project
|
|
82
|
+
gcp_region: us-central1
|
|
83
|
+
# Auth: GOOGLE_APPLICATION_CREDENTIALS or gcloud auth application-default login
|
|
84
|
+
"""
|
|
85
|
+
|
|
86
|
+
def __init__(
|
|
87
|
+
self,
|
|
88
|
+
provider: str = "anthropic",
|
|
89
|
+
model: Optional[str] = None,
|
|
90
|
+
api_key: Optional[str] = None,
|
|
91
|
+
token_tracker: Optional["TokenUsageTracker"] = None,
|
|
92
|
+
# Cloud provider config
|
|
93
|
+
aws_region: Optional[str] = None,
|
|
94
|
+
azure_endpoint: Optional[str] = None,
|
|
95
|
+
azure_deployment: Optional[str] = None,
|
|
96
|
+
azure_api_version: Optional[str] = None,
|
|
97
|
+
gcp_project: Optional[str] = None,
|
|
98
|
+
gcp_region: Optional[str] = None,
|
|
99
|
+
):
|
|
100
|
+
"""Initialize the LLM client.
|
|
101
|
+
|
|
102
|
+
Args:
|
|
103
|
+
provider: One of ``anthropic``, ``openai``, ``ollama``, ``google``, ``aws``, ``azure``, ``gcp``.
|
|
104
|
+
model: Model identifier (see defaults per provider below).
|
|
105
|
+
api_key: API key. If None, resolved from environment variables.
|
|
106
|
+
token_tracker: Optional :class:`~corbell.core.token_tracker.TokenUsageTracker`.
|
|
107
|
+
Each API call records its token usage here.
|
|
108
|
+
aws_region: AWS region for Bedrock (default: ``us-east-1``).
|
|
109
|
+
azure_endpoint: Azure OpenAI resource endpoint URL.
|
|
110
|
+
azure_deployment: Azure OpenAI deployment name.
|
|
111
|
+
azure_api_version: Azure OpenAI API version (default: ``2024-02-01``).
|
|
112
|
+
gcp_project: GCP project ID for Vertex AI.
|
|
113
|
+
gcp_region: GCP region for Vertex AI (default: ``us-central1``).
|
|
114
|
+
"""
|
|
115
|
+
self.provider = provider.lower()
|
|
116
|
+
self._api_key = api_key or self._resolve_key()
|
|
117
|
+
self.token_tracker = token_tracker
|
|
118
|
+
|
|
119
|
+
# Google multi-key support — parsed eagerly so is_configured is accurate
|
|
120
|
+
if self.provider == "google":
|
|
121
|
+
raw = self._api_key or ""
|
|
122
|
+
self._google_keys: list[str] = [k.strip() for k in raw.split(",") if k.strip()]
|
|
123
|
+
self._google_key_index: int = random.randrange(len(self._google_keys))
|
|
124
|
+
else:
|
|
125
|
+
self._google_keys = []
|
|
126
|
+
self._google_key_index = 0
|
|
127
|
+
|
|
128
|
+
# Cloud config
|
|
129
|
+
self.aws_region = aws_region or os.getenv("AWS_REGION", "us-east-1")
|
|
130
|
+
self.azure_endpoint = azure_endpoint or os.getenv("AZURE_OPENAI_ENDPOINT", "")
|
|
131
|
+
self.azure_deployment = azure_deployment or os.getenv("AZURE_OPENAI_DEPLOYMENT", "")
|
|
132
|
+
self.azure_api_version = azure_api_version or os.getenv("AZURE_OPENAI_API_VERSION", "2024-02-01")
|
|
133
|
+
self.gcp_project = gcp_project or os.getenv("GCP_PROJECT", os.getenv("GOOGLE_CLOUD_PROJECT", ""))
|
|
134
|
+
self.gcp_region = gcp_region or os.getenv("GCP_REGION", "us-central1")
|
|
135
|
+
|
|
136
|
+
_defaults = {
|
|
137
|
+
"anthropic": "claude-sonnet-4-5",
|
|
138
|
+
"openai": "gpt-4o",
|
|
139
|
+
"ollama": "llama3",
|
|
140
|
+
"google": "gemini-2.5-flash",
|
|
141
|
+
# Cloud defaults — Claude Sonnet 4.5 on Bedrock / Vertex
|
|
142
|
+
"aws": "us.anthropic.claude-sonnet-4-20250514-v1:0",
|
|
143
|
+
"azure": "gpt-4o",
|
|
144
|
+
"gcp": "claude-sonnet-4-5@20250514",
|
|
145
|
+
}
|
|
146
|
+
self.model = model or _defaults.get(self.provider, "claude-sonnet-4-5")
|
|
147
|
+
|
|
148
|
+
# ------------------------------------------------------------------ #
|
|
149
|
+
# Public API #
|
|
150
|
+
# ------------------------------------------------------------------ #
|
|
151
|
+
|
|
152
|
+
def call(
|
|
153
|
+
self,
|
|
154
|
+
system_prompt: str,
|
|
155
|
+
user_prompt: str,
|
|
156
|
+
max_tokens: int = 8000,
|
|
157
|
+
temperature: float = 0.1,
|
|
158
|
+
request_type: Optional[str] = None,
|
|
159
|
+
) -> str:
|
|
160
|
+
"""Call the configured LLM provider.
|
|
161
|
+
|
|
162
|
+
Args:
|
|
163
|
+
system_prompt: System / persona prompt.
|
|
164
|
+
user_prompt: User message / context.
|
|
165
|
+
max_tokens: Max tokens in the response.
|
|
166
|
+
temperature: Sampling temperature.
|
|
167
|
+
request_type: Label for token tracking (e.g. ``spec_generation``).
|
|
168
|
+
|
|
169
|
+
Returns:
|
|
170
|
+
Text response string. Falls back to structured template if no credentials.
|
|
171
|
+
"""
|
|
172
|
+
rt = request_type or "call"
|
|
173
|
+
|
|
174
|
+
provider_map = {
|
|
175
|
+
"anthropic": lambda: self._call_anthropic(system_prompt, user_prompt, max_tokens, temperature, rt),
|
|
176
|
+
"openai": lambda: self._call_openai(system_prompt, user_prompt, max_tokens, temperature, rt),
|
|
177
|
+
"ollama": lambda: self._call_ollama(system_prompt, user_prompt, max_tokens),
|
|
178
|
+
"google": lambda: self._call_google_ai(system_prompt, user_prompt, max_tokens, temperature, rt),
|
|
179
|
+
"aws": lambda: self._call_aws_bedrock(system_prompt, user_prompt, max_tokens, temperature, rt),
|
|
180
|
+
"azure": lambda: self._call_azure_openai(system_prompt, user_prompt, max_tokens, temperature, rt),
|
|
181
|
+
"gcp": lambda: self._call_gcp_vertex(system_prompt, user_prompt, max_tokens, temperature, rt),
|
|
182
|
+
}
|
|
183
|
+
|
|
184
|
+
if self.provider not in provider_map:
|
|
185
|
+
return self._fallback_response(system_prompt, user_prompt)
|
|
186
|
+
|
|
187
|
+
if not self.is_configured:
|
|
188
|
+
return self._fallback_response(system_prompt, user_prompt)
|
|
189
|
+
|
|
190
|
+
try:
|
|
191
|
+
return provider_map[self.provider]()
|
|
192
|
+
except Exception as e:
|
|
193
|
+
print(f"⚠️ LLM call failed ({self.provider}): {e}")
|
|
194
|
+
return self._fallback_response(system_prompt, user_prompt)
|
|
195
|
+
|
|
196
|
+
@property
|
|
197
|
+
def is_configured(self) -> bool:
|
|
198
|
+
"""True if credentials are available for the configured provider."""
|
|
199
|
+
if self.provider == "ollama":
|
|
200
|
+
return True
|
|
201
|
+
if self.provider == "aws":
|
|
202
|
+
# Long-term API key (BEDROCK_API_KEY) takes priority
|
|
203
|
+
if os.getenv("BEDROCK_API_KEY") or self._api_key:
|
|
204
|
+
return True
|
|
205
|
+
# Fall back: boto3 credential chain
|
|
206
|
+
return bool(
|
|
207
|
+
os.getenv("AWS_ACCESS_KEY_ID")
|
|
208
|
+
or os.getenv("AWS_PROFILE")
|
|
209
|
+
or os.path.exists(os.path.expanduser("~/.aws/credentials"))
|
|
210
|
+
)
|
|
211
|
+
if self.provider == "azure":
|
|
212
|
+
return bool(self._api_key and self.azure_endpoint)
|
|
213
|
+
if self.provider == "gcp":
|
|
214
|
+
return bool(
|
|
215
|
+
os.getenv("GOOGLE_APPLICATION_CREDENTIALS")
|
|
216
|
+
or os.getenv("GOOGLE_CLOUD_PROJECT")
|
|
217
|
+
or self.gcp_project
|
|
218
|
+
)
|
|
219
|
+
if self.provider == "google":
|
|
220
|
+
return bool(self._google_keys)
|
|
221
|
+
return bool(self._api_key)
|
|
222
|
+
|
|
223
|
+
@property
|
|
224
|
+
def provider_display(self) -> str:
|
|
225
|
+
"""Human-readable provider description."""
|
|
226
|
+
labels = {
|
|
227
|
+
"anthropic": f"Anthropic ({self.model})",
|
|
228
|
+
"openai": f"OpenAI ({self.model})",
|
|
229
|
+
"ollama": f"Ollama/{self.model} (local)",
|
|
230
|
+
"google": f"Google AI ({self.model})",
|
|
231
|
+
"aws": f"AWS Bedrock/{self.model} @ {self.aws_region}",
|
|
232
|
+
"azure": f"Azure OpenAI/{self.model} ({self.azure_deployment or 'default'})",
|
|
233
|
+
"gcp": f"GCP Vertex AI/{self.model} @ {self.gcp_region}",
|
|
234
|
+
}
|
|
235
|
+
return labels.get(self.provider, self.provider)
|
|
236
|
+
|
|
237
|
+
# ------------------------------------------------------------------ #
|
|
238
|
+
# Provider implementations — local #
|
|
239
|
+
# ------------------------------------------------------------------ #
|
|
240
|
+
|
|
241
|
+
def _call_anthropic(
|
|
242
|
+
self, system: str, user: str, max_tokens: int, temperature: float,
|
|
243
|
+
request_type: str = "call",
|
|
244
|
+
) -> str:
|
|
245
|
+
try:
|
|
246
|
+
import anthropic
|
|
247
|
+
except ImportError:
|
|
248
|
+
raise ImportError("pip install corbell[anthropic]")
|
|
249
|
+
|
|
250
|
+
client = anthropic.Anthropic(api_key=self._api_key)
|
|
251
|
+
msg = client.messages.create(
|
|
252
|
+
model=self.model,
|
|
253
|
+
max_tokens=max_tokens,
|
|
254
|
+
temperature=temperature,
|
|
255
|
+
system=system,
|
|
256
|
+
messages=[{"role": "user", "content": user}],
|
|
257
|
+
)
|
|
258
|
+
|
|
259
|
+
if self.token_tracker and hasattr(msg, "usage"):
|
|
260
|
+
self.token_tracker.record(request_type, self.model, msg.usage.input_tokens, msg.usage.output_tokens)
|
|
261
|
+
|
|
262
|
+
return msg.content[0].text
|
|
263
|
+
|
|
264
|
+
def _call_openai(
|
|
265
|
+
self, system: str, user: str, max_tokens: int, temperature: float,
|
|
266
|
+
request_type: str = "call",
|
|
267
|
+
) -> str:
|
|
268
|
+
try:
|
|
269
|
+
import openai
|
|
270
|
+
except ImportError:
|
|
271
|
+
raise ImportError("pip install corbell[openai]")
|
|
272
|
+
|
|
273
|
+
client = openai.OpenAI(api_key=self._api_key)
|
|
274
|
+
resp = client.chat.completions.create(
|
|
275
|
+
model=self.model,
|
|
276
|
+
max_tokens=max_tokens,
|
|
277
|
+
temperature=temperature,
|
|
278
|
+
messages=[
|
|
279
|
+
{"role": "system", "content": system},
|
|
280
|
+
{"role": "user", "content": user},
|
|
281
|
+
],
|
|
282
|
+
)
|
|
283
|
+
|
|
284
|
+
if self.token_tracker and resp.usage:
|
|
285
|
+
self.token_tracker.record(
|
|
286
|
+
request_type, self.model,
|
|
287
|
+
resp.usage.prompt_tokens, resp.usage.completion_tokens,
|
|
288
|
+
)
|
|
289
|
+
|
|
290
|
+
return resp.choices[0].message.content or ""
|
|
291
|
+
|
|
292
|
+
def _call_ollama(self, system: str, user: str, max_tokens: int) -> str:
|
|
293
|
+
"""Call a local Ollama instance (no token tracking — free local model)."""
|
|
294
|
+
import urllib.request
|
|
295
|
+
|
|
296
|
+
payload = json.dumps({
|
|
297
|
+
"model": self.model,
|
|
298
|
+
"messages": [
|
|
299
|
+
{"role": "system", "content": system},
|
|
300
|
+
{"role": "user", "content": user},
|
|
301
|
+
],
|
|
302
|
+
"stream": False,
|
|
303
|
+
"options": {"num_predict": max_tokens},
|
|
304
|
+
}).encode()
|
|
305
|
+
|
|
306
|
+
req = urllib.request.Request(
|
|
307
|
+
"http://localhost:11434/api/chat",
|
|
308
|
+
data=payload,
|
|
309
|
+
headers={"Content-Type": "application/json"},
|
|
310
|
+
)
|
|
311
|
+
with urllib.request.urlopen(req, timeout=120) as resp:
|
|
312
|
+
data = json.loads(resp.read())
|
|
313
|
+
return data.get("message", {}).get("content", "")
|
|
314
|
+
|
|
315
|
+
def _call_google_ai(
|
|
316
|
+
self, system: str, user: str, max_tokens: int, temperature: float,
|
|
317
|
+
request_type: str = "call",
|
|
318
|
+
) -> str:
|
|
319
|
+
"""Call Google AI (Gemini) models via the google-genai SDK.
|
|
320
|
+
|
|
321
|
+
Rotates across all configured keys on auth/quota failures.
|
|
322
|
+
Raises RuntimeError only when all keys are exhausted.
|
|
323
|
+
|
|
324
|
+
Requires ``pip install corbell[google]`` and ``GOOGLE_API_KEY``.
|
|
325
|
+
"""
|
|
326
|
+
try:
|
|
327
|
+
from google import genai
|
|
328
|
+
from google.genai import types
|
|
329
|
+
except ImportError:
|
|
330
|
+
raise ImportError("pip install corbell[google]")
|
|
331
|
+
|
|
332
|
+
start = self._google_key_index
|
|
333
|
+
errors: list[str] = []
|
|
334
|
+
for i in range(len(self._google_keys)):
|
|
335
|
+
idx = (start + i) % len(self._google_keys)
|
|
336
|
+
key = self._google_keys[idx]
|
|
337
|
+
try:
|
|
338
|
+
client = genai.Client(api_key=key)
|
|
339
|
+
response = client.models.generate_content(
|
|
340
|
+
model=self.model,
|
|
341
|
+
contents=user,
|
|
342
|
+
config=types.GenerateContentConfig(
|
|
343
|
+
system_instruction=system,
|
|
344
|
+
max_output_tokens=max_tokens,
|
|
345
|
+
temperature=temperature,
|
|
346
|
+
),
|
|
347
|
+
)
|
|
348
|
+
self._google_key_index = (idx + 1) % len(self._google_keys)
|
|
349
|
+
if self.token_tracker and response.usage_metadata:
|
|
350
|
+
self.token_tracker.record(
|
|
351
|
+
request_type, self.model,
|
|
352
|
+
response.usage_metadata.prompt_token_count or 0,
|
|
353
|
+
response.usage_metadata.candidates_token_count or 0,
|
|
354
|
+
)
|
|
355
|
+
return response.text
|
|
356
|
+
except Exception as e:
|
|
357
|
+
if self._is_google_key_error(e):
|
|
358
|
+
errors.append(f"key[{idx}]: {e}")
|
|
359
|
+
continue
|
|
360
|
+
raise
|
|
361
|
+
|
|
362
|
+
raise RuntimeError(
|
|
363
|
+
f"All {len(self._google_keys)} Google API key(s) failed:\n"
|
|
364
|
+
+ "\n".join(errors)
|
|
365
|
+
)
|
|
366
|
+
|
|
367
|
+
# ------------------------------------------------------------------ #
|
|
368
|
+
# Provider implementations — cloud #
|
|
369
|
+
# ------------------------------------------------------------------ #
|
|
370
|
+
|
|
371
|
+
def _call_aws_bedrock(
|
|
372
|
+
self, system: str, user: str, max_tokens: int, temperature: float,
|
|
373
|
+
request_type: str = "call",
|
|
374
|
+
) -> str:
|
|
375
|
+
"""Call Anthropic Claude via AWS Bedrock.
|
|
376
|
+
|
|
377
|
+
**Auth option 1 — Long-term API key (simplest, recommended):**
|
|
378
|
+
Paste the key AWS gives you from the Bedrock console directly.
|
|
379
|
+
|
|
380
|
+
.. code-block:: bash
|
|
381
|
+
|
|
382
|
+
export BEDROCK_API_KEY=your-long-term-api-key
|
|
383
|
+
export AWS_REGION=us-east-1 # optional, default: us-east-1
|
|
384
|
+
|
|
385
|
+
Or in workspace.yaml:
|
|
386
|
+
|
|
387
|
+
.. code-block:: yaml
|
|
388
|
+
|
|
389
|
+
llm:
|
|
390
|
+
provider: aws
|
|
391
|
+
model: us.anthropic.claude-sonnet-4-20250514-v1:0
|
|
392
|
+
api_key: ${BEDROCK_API_KEY}
|
|
393
|
+
aws_region: us-east-1
|
|
394
|
+
|
|
395
|
+
**Auth option 2 — IAM credential chain (boto3):**
|
|
396
|
+
- Environment: ``AWS_ACCESS_KEY_ID`` + ``AWS_SECRET_ACCESS_KEY``
|
|
397
|
+
- Profile: ``aws configure`` or ``AWS_PROFILE``
|
|
398
|
+
- Instance metadata (EC2/ECS/Lambda)
|
|
399
|
+
"""
|
|
400
|
+
# --- Long-term Bearer key path (simplest for users) ---
|
|
401
|
+
bearer_key = os.getenv("BEDROCK_API_KEY") or self._api_key
|
|
402
|
+
region = self.aws_region or os.getenv("AWS_REGION", "us-east-1")
|
|
403
|
+
endpoint_url = f"https://bedrock-runtime.{region}.amazonaws.com/model/{self.model}/invoke"
|
|
404
|
+
|
|
405
|
+
if bearer_key:
|
|
406
|
+
import urllib.request
|
|
407
|
+
import urllib.error
|
|
408
|
+
|
|
409
|
+
payload = json.dumps({
|
|
410
|
+
"anthropic_version": "bedrock-2023-05-31",
|
|
411
|
+
"max_tokens": max_tokens,
|
|
412
|
+
"temperature": temperature,
|
|
413
|
+
"system": system,
|
|
414
|
+
"messages": [{"role": "user", "content": user}],
|
|
415
|
+
"top_p": 0.9,
|
|
416
|
+
}).encode()
|
|
417
|
+
|
|
418
|
+
req = urllib.request.Request(
|
|
419
|
+
endpoint_url,
|
|
420
|
+
data=payload,
|
|
421
|
+
headers={
|
|
422
|
+
"Authorization": f"Bearer {bearer_key}",
|
|
423
|
+
"Content-Type": "application/json",
|
|
424
|
+
},
|
|
425
|
+
)
|
|
426
|
+
with urllib.request.urlopen(req, timeout=180) as resp:
|
|
427
|
+
result = json.loads(resp.read())
|
|
428
|
+
|
|
429
|
+
if self.token_tracker:
|
|
430
|
+
usage = result.get("usage", {})
|
|
431
|
+
self.token_tracker.record(
|
|
432
|
+
request_type, self.model,
|
|
433
|
+
usage.get("input_tokens", 0),
|
|
434
|
+
usage.get("output_tokens", 0),
|
|
435
|
+
)
|
|
436
|
+
content = result.get("content", [])
|
|
437
|
+
if content:
|
|
438
|
+
return content[0]["text"]
|
|
439
|
+
raise ValueError(f"Unexpected Bedrock response: {result}")
|
|
440
|
+
|
|
441
|
+
# --- boto3 IAM credential chain fallback ---
|
|
442
|
+
try:
|
|
443
|
+
import boto3
|
|
444
|
+
except ImportError:
|
|
445
|
+
raise ImportError(
|
|
446
|
+
"pip install corbell[aws]\n"
|
|
447
|
+
"Then either:\n"
|
|
448
|
+
" Option 1 (simpler): set BEDROCK_API_KEY=<your AWS Bedrock key>\n"
|
|
449
|
+
" Option 2 (IAM): aws configure (or set AWS_ACCESS_KEY_ID/SECRET)"
|
|
450
|
+
)
|
|
451
|
+
|
|
452
|
+
client = boto3.client("bedrock-runtime", region_name=region)
|
|
453
|
+
body = json.dumps({
|
|
454
|
+
"anthropic_version": "bedrock-2023-05-31",
|
|
455
|
+
"max_tokens": max_tokens,
|
|
456
|
+
"temperature": temperature,
|
|
457
|
+
"system": system,
|
|
458
|
+
"messages": [{"role": "user", "content": user}],
|
|
459
|
+
})
|
|
460
|
+
|
|
461
|
+
resp = client.invoke_model(modelId=self.model, body=body)
|
|
462
|
+
result = json.loads(resp["body"].read())
|
|
463
|
+
|
|
464
|
+
if self.token_tracker:
|
|
465
|
+
usage = result.get("usage", {})
|
|
466
|
+
self.token_tracker.record(
|
|
467
|
+
request_type, self.model,
|
|
468
|
+
usage.get("input_tokens", 0),
|
|
469
|
+
usage.get("output_tokens", 0),
|
|
470
|
+
)
|
|
471
|
+
|
|
472
|
+
content = result.get("content", [])
|
|
473
|
+
if content:
|
|
474
|
+
return content[0]["text"]
|
|
475
|
+
raise ValueError(f"Unexpected Bedrock response: {result}")
|
|
476
|
+
|
|
477
|
+
def _call_azure_openai(
|
|
478
|
+
self, system: str, user: str, max_tokens: int, temperature: float,
|
|
479
|
+
request_type: str = "call",
|
|
480
|
+
) -> str:
|
|
481
|
+
"""Call GPT-4 models via Azure OpenAI Service.
|
|
482
|
+
|
|
483
|
+
Requires:
|
|
484
|
+
- AZURE_OPENAI_API_KEY (or api_key in workspace.yaml)
|
|
485
|
+
- AZURE_OPENAI_ENDPOINT (e.g. https://my-resource.openai.azure.com/)
|
|
486
|
+
- AZURE_OPENAI_DEPLOYMENT (your deployment name, e.g. gpt-4o-prod)
|
|
487
|
+
|
|
488
|
+
Set these in your .env or workspace.yaml llm block.
|
|
489
|
+
"""
|
|
490
|
+
try:
|
|
491
|
+
import openai
|
|
492
|
+
except ImportError:
|
|
493
|
+
raise ImportError("pip install corbell[openai]")
|
|
494
|
+
|
|
495
|
+
deployment = self.azure_deployment or self.model
|
|
496
|
+
client = openai.AzureOpenAI(
|
|
497
|
+
api_key=self._api_key,
|
|
498
|
+
azure_endpoint=self.azure_endpoint,
|
|
499
|
+
azure_deployment=deployment,
|
|
500
|
+
api_version=self.azure_api_version,
|
|
501
|
+
)
|
|
502
|
+
|
|
503
|
+
resp = client.chat.completions.create(
|
|
504
|
+
model=deployment,
|
|
505
|
+
max_tokens=max_tokens,
|
|
506
|
+
temperature=temperature,
|
|
507
|
+
messages=[
|
|
508
|
+
{"role": "system", "content": system},
|
|
509
|
+
{"role": "user", "content": user},
|
|
510
|
+
],
|
|
511
|
+
)
|
|
512
|
+
|
|
513
|
+
if self.token_tracker and resp.usage:
|
|
514
|
+
self.token_tracker.record(
|
|
515
|
+
request_type, f"azure/{deployment}",
|
|
516
|
+
resp.usage.prompt_tokens, resp.usage.completion_tokens,
|
|
517
|
+
)
|
|
518
|
+
|
|
519
|
+
return resp.choices[0].message.content or ""
|
|
520
|
+
|
|
521
|
+
def _call_gcp_vertex(
|
|
522
|
+
self, system: str, user: str, max_tokens: int, temperature: float,
|
|
523
|
+
request_type: str = "call",
|
|
524
|
+
) -> str:
|
|
525
|
+
"""Call Anthropic Claude via GCP Vertex AI.
|
|
526
|
+
|
|
527
|
+
Auth options (pick one):
|
|
528
|
+
1. Application Default Credentials: ``gcloud auth application-default login``
|
|
529
|
+
2. Service account: set GOOGLE_APPLICATION_CREDENTIALS=/path/to/key.json
|
|
530
|
+
|
|
531
|
+
Set GCP_PROJECT + GCP_REGION in .env or workspace.yaml.
|
|
532
|
+
|
|
533
|
+
Requires: ``pip install "google-cloud-aiplatform>=1.38" anthropic[vertex]``
|
|
534
|
+
"""
|
|
535
|
+
try:
|
|
536
|
+
from anthropic import AnthropicVertex # type: ignore[attr-defined]
|
|
537
|
+
except (ImportError, AttributeError):
|
|
538
|
+
raise ImportError(
|
|
539
|
+
"pip install anthropic[vertex] google-cloud-aiplatform\n"
|
|
540
|
+
"Then authenticate: gcloud auth application-default login\n"
|
|
541
|
+
"Or set GOOGLE_APPLICATION_CREDENTIALS=/path/to/service-account-key.json"
|
|
542
|
+
)
|
|
543
|
+
|
|
544
|
+
client = AnthropicVertex(
|
|
545
|
+
project_id=self.gcp_project,
|
|
546
|
+
region=self.gcp_region,
|
|
547
|
+
)
|
|
548
|
+
|
|
549
|
+
msg = client.messages.create(
|
|
550
|
+
model=self.model,
|
|
551
|
+
max_tokens=max_tokens,
|
|
552
|
+
temperature=temperature,
|
|
553
|
+
system=system,
|
|
554
|
+
messages=[{"role": "user", "content": user}],
|
|
555
|
+
)
|
|
556
|
+
|
|
557
|
+
if self.token_tracker and hasattr(msg, "usage"):
|
|
558
|
+
self.token_tracker.record(
|
|
559
|
+
request_type, f"gcp/{self.model}",
|
|
560
|
+
msg.usage.input_tokens, msg.usage.output_tokens,
|
|
561
|
+
)
|
|
562
|
+
|
|
563
|
+
return msg.content[0].text
|
|
564
|
+
|
|
565
|
+
# ------------------------------------------------------------------ #
|
|
566
|
+
# Helpers #
|
|
567
|
+
# ------------------------------------------------------------------ #
|
|
568
|
+
|
|
569
|
+
@staticmethod
|
|
570
|
+
def _is_google_key_error(e: Exception) -> bool:
|
|
571
|
+
"""Return True when a Google ClientError is caused by the key, not the request."""
|
|
572
|
+
code = getattr(e, "code", None)
|
|
573
|
+
if code in (401, 403, 429):
|
|
574
|
+
return True
|
|
575
|
+
if code == 400:
|
|
576
|
+
msg = (getattr(e, "message", None) or str(e)).lower()
|
|
577
|
+
return "api key" in msg
|
|
578
|
+
return False
|
|
579
|
+
|
|
580
|
+
def _resolve_key(self) -> Optional[str]:
|
|
581
|
+
env_map = {
|
|
582
|
+
"anthropic": ["ANTHROPIC_API_KEY", "CORBELL_LLM_API_KEY"],
|
|
583
|
+
"openai": ["OPENAI_API_KEY", "CORBELL_LLM_API_KEY"],
|
|
584
|
+
"azure": ["AZURE_OPENAI_API_KEY", "CORBELL_LLM_API_KEY"],
|
|
585
|
+
"google": ["GOOGLE_API_KEY", "CORBELL_LLM_API_KEY"],
|
|
586
|
+
"ollama": [],
|
|
587
|
+
"aws": [], # Uses boto3 credential chain
|
|
588
|
+
"gcp": [], # Uses Google ADC
|
|
589
|
+
}
|
|
590
|
+
for var in env_map.get(self.provider, ["CORBELL_LLM_API_KEY"]):
|
|
591
|
+
val = os.environ.get(var)
|
|
592
|
+
if val:
|
|
593
|
+
return val
|
|
594
|
+
return None
|
|
595
|
+
|
|
596
|
+
def _fallback_response(self, system: str, user: str) -> str:
|
|
597
|
+
"""Return a structured template when no LLM credentials are available."""
|
|
598
|
+
if "design document" in system.lower() or "technical design" in system.lower():
|
|
599
|
+
return _MOCK_DESIGN_DOC
|
|
600
|
+
if "design decisions" in system.lower() or "extract" in system.lower():
|
|
601
|
+
return "[]"
|
|
602
|
+
if "pattern" in system.lower():
|
|
603
|
+
return "{}"
|
|
604
|
+
if any(kw in system.lower() for kw in ("search", "keywords", "queries")):
|
|
605
|
+
import re
|
|
606
|
+
sentences = [s.strip() for s in re.split(r'[.\n]', user) if len(s.strip()) > 30]
|
|
607
|
+
return "\n".join(sentences[:3]) if sentences else user[:200]
|
|
608
|
+
return (
|
|
609
|
+
"⚠️ No LLM credentials configured.\n"
|
|
610
|
+
"\n"
|
|
611
|
+
"Quick setup — pick your provider:\n"
|
|
612
|
+
"\n"
|
|
613
|
+
" Anthropic: export ANTHROPIC_API_KEY=sk-ant-...\n"
|
|
614
|
+
" OpenAI: export OPENAI_API_KEY=sk-...\n"
|
|
615
|
+
" Google AI: export GOOGLE_API_KEY=AIza...\n"
|
|
616
|
+
" AWS Bedrock: export BEDROCK_API_KEY=<your-long-term-key> AWS_REGION=us-east-1\n"
|
|
617
|
+
" (or IAM): export AWS_ACCESS_KEY_ID=... AWS_SECRET_ACCESS_KEY=...\n"
|
|
618
|
+
" Azure: export AZURE_OPENAI_API_KEY=... AZURE_OPENAI_ENDPOINT=https://...\n"
|
|
619
|
+
" export AZURE_OPENAI_DEPLOYMENT=my-gpt4o\n"
|
|
620
|
+
" GCP Vertex: gcloud auth application-default login\n"
|
|
621
|
+
" export GCP_PROJECT=my-project GCP_REGION=us-central1\n"
|
|
622
|
+
"\n"
|
|
623
|
+
"Update corbell-data/workspace.yaml llm.provider accordingly, then re-run."
|
|
624
|
+
)
|
|
625
|
+
|
|
626
|
+
|
|
627
|
+
_MOCK_DESIGN_DOC = """\
|
|
628
|
+
# Technical Design Document
|
|
629
|
+
|
|
630
|
+
> ⚠️ **Template mode**: No LLM credentials configured.
|
|
631
|
+
>
|
|
632
|
+
> Quick setup options:
|
|
633
|
+
> - Anthropic: `export ANTHROPIC_API_KEY=sk-ant-...`
|
|
634
|
+
> - OpenAI: `export OPENAI_API_KEY=sk-...`
|
|
635
|
+
> - AWS Bedrock: `export AWS_ACCESS_KEY_ID=... AWS_SECRET_ACCESS_KEY=... AWS_REGION=us-east-1`
|
|
636
|
+
> - Azure: `export AZURE_OPENAI_API_KEY=... AZURE_OPENAI_ENDPOINT=https://...`
|
|
637
|
+
> - GCP Vertex: `gcloud auth application-default login && export GCP_PROJECT=...`
|
|
638
|
+
>
|
|
639
|
+
> See README.md → LLM Providers for full instructions.
|
|
640
|
+
|
|
641
|
+
## Context
|
|
642
|
+
|
|
643
|
+
<!-- Describe WHY this feature is being built. -->
|
|
644
|
+
|
|
645
|
+
## Current Architecture
|
|
646
|
+
|
|
647
|
+
<!-- CORBELL_GRAPH_START -->
|
|
648
|
+
<!-- Current service graph will be inserted here by corbell. -->
|
|
649
|
+
<!-- CORBELL_GRAPH_END -->
|
|
650
|
+
|
|
651
|
+
## Proposed Design
|
|
652
|
+
|
|
653
|
+
### Service Changes
|
|
654
|
+
|
|
655
|
+
<!-- What changes in each service. -->
|
|
656
|
+
|
|
657
|
+
### Data Flow
|
|
658
|
+
|
|
659
|
+
<!-- Sequence or description of how data moves. -->
|
|
660
|
+
|
|
661
|
+
### Failure Modes and Mitigations
|
|
662
|
+
|
|
663
|
+
<!-- What can go wrong, how each is handled. -->
|
|
664
|
+
|
|
665
|
+
## Reliability and Risk Constraints
|
|
666
|
+
|
|
667
|
+
<!-- CORBELL_CONSTRAINTS_START -->
|
|
668
|
+
<!-- CORBELL_CONSTRAINTS_END -->
|
|
669
|
+
|
|
670
|
+
## Rollout Plan
|
|
671
|
+
|
|
672
|
+
<!-- Phases, feature flags, rollback plan. -->
|
|
673
|
+
|
|
674
|
+
## Open Questions
|
|
675
|
+
|
|
676
|
+
<!-- Things not yet decided. -->
|
|
677
|
+
"""
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
"""Init file for MCP module"""
|