switchforge 1.0.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.
- switchforge-1.0.0/.gitignore +36 -0
- switchforge-1.0.0/PKG-INFO +65 -0
- switchforge-1.0.0/README.md +31 -0
- switchforge-1.0.0/forge-core.spec +108 -0
- switchforge-1.0.0/forge_core/__init__.py +3 -0
- switchforge-1.0.0/forge_core/__main__.py +6 -0
- switchforge-1.0.0/forge_core/ai/__init__.py +0 -0
- switchforge-1.0.0/forge_core/ai/prompts.py +87 -0
- switchforge-1.0.0/forge_core/ai/provider.py +121 -0
- switchforge-1.0.0/forge_core/ai/structured.py +68 -0
- switchforge-1.0.0/forge_core/auth.py +72 -0
- switchforge-1.0.0/forge_core/cli.py +209 -0
- switchforge-1.0.0/forge_core/config.py +118 -0
- switchforge-1.0.0/forge_core/core/__init__.py +0 -0
- switchforge-1.0.0/forge_core/core/agent_manager.py +111 -0
- switchforge-1.0.0/forge_core/core/coverage.py +241 -0
- switchforge-1.0.0/forge_core/core/file_manager.py +104 -0
- switchforge-1.0.0/forge_core/models/__init__.py +0 -0
- switchforge-1.0.0/forge_core/models/config.py +115 -0
- switchforge-1.0.0/forge_core/models/dto.py +58 -0
- switchforge-1.0.0/forge_core/models/project.py +72 -0
- switchforge-1.0.0/forge_core/models/test_result.py +93 -0
- switchforge-1.0.0/forge_core/orchestrator.py +233 -0
- switchforge-1.0.0/forge_core/phases/__init__.py +0 -0
- switchforge-1.0.0/forge_core/phases/analyze_project.py +149 -0
- switchforge-1.0.0/forge_core/phases/audit_tests.py +35 -0
- switchforge-1.0.0/forge_core/phases/compile_fix.py +79 -0
- switchforge-1.0.0/forge_core/phases/coverage_report.py +19 -0
- switchforge-1.0.0/forge_core/phases/detect_stack.py +66 -0
- switchforge-1.0.0/forge_core/phases/exclusion_scan.py +74 -0
- switchforge-1.0.0/forge_core/phases/fix_broken.py +83 -0
- switchforge-1.0.0/forge_core/phases/generate_tests.py +218 -0
- switchforge-1.0.0/forge_core/phases/journey_mapping.py +120 -0
- switchforge-1.0.0/forge_core/phases/self_learn.py +93 -0
- switchforge-1.0.0/forge_core/utils/__init__.py +0 -0
- switchforge-1.0.0/forge_core/utils/logger.py +89 -0
- switchforge-1.0.0/forge_core/utils/reporter.py +70 -0
- switchforge-1.0.0/forge_core/utils/shell.py +93 -0
- switchforge-1.0.0/forge_core/utils/tokens.py +67 -0
- switchforge-1.0.0/pyproject.toml +66 -0
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
# OS
|
|
2
|
+
.DS_Store
|
|
3
|
+
Thumbs.db
|
|
4
|
+
|
|
5
|
+
# IDE
|
|
6
|
+
.idea/
|
|
7
|
+
.vscode/
|
|
8
|
+
*.swp
|
|
9
|
+
*.swo
|
|
10
|
+
|
|
11
|
+
# Build outputs
|
|
12
|
+
build/
|
|
13
|
+
dist/
|
|
14
|
+
out/
|
|
15
|
+
target/
|
|
16
|
+
|
|
17
|
+
# Dependencies
|
|
18
|
+
node_modules/
|
|
19
|
+
__pycache__/
|
|
20
|
+
*.pyc
|
|
21
|
+
.venv/
|
|
22
|
+
venv/
|
|
23
|
+
|
|
24
|
+
# Environment
|
|
25
|
+
.env
|
|
26
|
+
.env.local
|
|
27
|
+
.env.*.local
|
|
28
|
+
|
|
29
|
+
# Logs
|
|
30
|
+
*.log
|
|
31
|
+
npm-debug.log*
|
|
32
|
+
|
|
33
|
+
# Coverage
|
|
34
|
+
coverage/
|
|
35
|
+
htmlcov/
|
|
36
|
+
.coverage
|
|
@@ -0,0 +1,65 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: switchforge
|
|
3
|
+
Version: 1.0.0
|
|
4
|
+
Summary: AI-powered backend test generation engine
|
|
5
|
+
Project-URL: Homepage, https://theswitchcompany.online/products/forge/core
|
|
6
|
+
Project-URL: Repository, https://github.com/switchcompany/forge-core
|
|
7
|
+
Project-URL: Documentation, https://theswitchcompany.online/docs/forge-core
|
|
8
|
+
Author-email: TheSwitchCompany <hello@theswitchcompany.online>
|
|
9
|
+
License-Expression: MIT
|
|
10
|
+
Keywords: ai,backend,code-generation,testing,unit-tests
|
|
11
|
+
Classifier: Development Status :: 4 - Beta
|
|
12
|
+
Classifier: Intended Audience :: Developers
|
|
13
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
14
|
+
Classifier: Programming Language :: Python :: 3
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
16
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
17
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
18
|
+
Classifier: Topic :: Software Development :: Testing
|
|
19
|
+
Requires-Python: >=3.10
|
|
20
|
+
Requires-Dist: httpx>=0.27.0
|
|
21
|
+
Requires-Dist: instructor>=1.3.0
|
|
22
|
+
Requires-Dist: litellm>=1.40.0
|
|
23
|
+
Requires-Dist: pydantic>=2.7.0
|
|
24
|
+
Requires-Dist: pyyaml>=6.0
|
|
25
|
+
Requires-Dist: rich>=13.7.0
|
|
26
|
+
Requires-Dist: tiktoken>=0.7.0
|
|
27
|
+
Requires-Dist: typer>=0.12.0
|
|
28
|
+
Provides-Extra: dev
|
|
29
|
+
Requires-Dist: pyinstaller>=6.0; extra == 'dev'
|
|
30
|
+
Requires-Dist: pytest-asyncio>=0.23; extra == 'dev'
|
|
31
|
+
Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
32
|
+
Requires-Dist: ruff>=0.5.0; extra == 'dev'
|
|
33
|
+
Description-Content-Type: text/markdown
|
|
34
|
+
|
|
35
|
+
# Switchforge (Forge Core CLI)
|
|
36
|
+
|
|
37
|
+
AI-powered backend test generation engine by [TheSwitchCompany](https://theswitchcompany.online).
|
|
38
|
+
|
|
39
|
+
## Install
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
pip install switchforge
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
## Quick Start
|
|
46
|
+
|
|
47
|
+
```bash
|
|
48
|
+
switchforge login --token YOUR_API_TOKEN
|
|
49
|
+
switchforge init
|
|
50
|
+
switchforge run . --target 90
|
|
51
|
+
```
|
|
52
|
+
|
|
53
|
+
### Windows Users
|
|
54
|
+
|
|
55
|
+
If `switchforge` is not recognized, use `python -m forge_core`:
|
|
56
|
+
|
|
57
|
+
```powershell
|
|
58
|
+
python -m forge_core login --token YOUR_API_TOKEN
|
|
59
|
+
python -m forge_core init
|
|
60
|
+
python -m forge_core run . --target 90
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
## Documentation
|
|
64
|
+
|
|
65
|
+
[theswitchcompany.online/docs/forge-core](https://theswitchcompany.online/docs/forge-core)
|
|
@@ -0,0 +1,31 @@
|
|
|
1
|
+
# Switchforge (Forge Core CLI)
|
|
2
|
+
|
|
3
|
+
AI-powered backend test generation engine by [TheSwitchCompany](https://theswitchcompany.online).
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install switchforge
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Quick Start
|
|
12
|
+
|
|
13
|
+
```bash
|
|
14
|
+
switchforge login --token YOUR_API_TOKEN
|
|
15
|
+
switchforge init
|
|
16
|
+
switchforge run . --target 90
|
|
17
|
+
```
|
|
18
|
+
|
|
19
|
+
### Windows Users
|
|
20
|
+
|
|
21
|
+
If `switchforge` is not recognized, use `python -m forge_core`:
|
|
22
|
+
|
|
23
|
+
```powershell
|
|
24
|
+
python -m forge_core login --token YOUR_API_TOKEN
|
|
25
|
+
python -m forge_core init
|
|
26
|
+
python -m forge_core run . --target 90
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
## Documentation
|
|
30
|
+
|
|
31
|
+
[theswitchcompany.online/docs/forge-core](https://theswitchcompany.online/docs/forge-core)
|
|
@@ -0,0 +1,108 @@
|
|
|
1
|
+
# -*- mode: python ; coding: utf-8 -*-
|
|
2
|
+
"""
|
|
3
|
+
PyInstaller spec for Forge Core CLI.
|
|
4
|
+
Compiles the Python engine into a single standalone binary.
|
|
5
|
+
|
|
6
|
+
Build:
|
|
7
|
+
pip install pyinstaller
|
|
8
|
+
pyinstaller forge-core.spec
|
|
9
|
+
|
|
10
|
+
Output: dist/forge-core (single binary, ~50-80MB)
|
|
11
|
+
"""
|
|
12
|
+
|
|
13
|
+
import sys
|
|
14
|
+
from pathlib import Path
|
|
15
|
+
|
|
16
|
+
block_cipher = None
|
|
17
|
+
|
|
18
|
+
# Collect all forge_core modules
|
|
19
|
+
a = Analysis(
|
|
20
|
+
['forge_core/cli.py'],
|
|
21
|
+
pathex=[],
|
|
22
|
+
binaries=[],
|
|
23
|
+
datas=[
|
|
24
|
+
# Include prompt templates if bundled
|
|
25
|
+
('../.github/prompts', 'prompts'),
|
|
26
|
+
('../knowledge-packs', 'knowledge-packs'),
|
|
27
|
+
],
|
|
28
|
+
hiddenimports=[
|
|
29
|
+
'forge_core',
|
|
30
|
+
'forge_core.orchestrator',
|
|
31
|
+
'forge_core.config',
|
|
32
|
+
'forge_core.auth',
|
|
33
|
+
'forge_core.ai.provider',
|
|
34
|
+
'forge_core.ai.prompts',
|
|
35
|
+
'forge_core.ai.structured',
|
|
36
|
+
'forge_core.phases.detect_stack',
|
|
37
|
+
'forge_core.phases.exclusion_scan',
|
|
38
|
+
'forge_core.phases.analyze_project',
|
|
39
|
+
'forge_core.phases.journey_mapping',
|
|
40
|
+
'forge_core.phases.audit_tests',
|
|
41
|
+
'forge_core.phases.fix_broken',
|
|
42
|
+
'forge_core.phases.generate_tests',
|
|
43
|
+
'forge_core.phases.compile_fix',
|
|
44
|
+
'forge_core.phases.coverage_report',
|
|
45
|
+
'forge_core.phases.self_learn',
|
|
46
|
+
'forge_core.core.file_manager',
|
|
47
|
+
'forge_core.core.coverage',
|
|
48
|
+
'forge_core.core.agent_manager',
|
|
49
|
+
'forge_core.models.config',
|
|
50
|
+
'forge_core.models.project',
|
|
51
|
+
'forge_core.models.dto',
|
|
52
|
+
'forge_core.models.test_result',
|
|
53
|
+
'forge_core.utils.logger',
|
|
54
|
+
'forge_core.utils.shell',
|
|
55
|
+
'forge_core.utils.tokens',
|
|
56
|
+
'forge_core.utils.reporter',
|
|
57
|
+
# Dependencies
|
|
58
|
+
'litellm',
|
|
59
|
+
'instructor',
|
|
60
|
+
'typer',
|
|
61
|
+
'rich',
|
|
62
|
+
'yaml',
|
|
63
|
+
'tiktoken',
|
|
64
|
+
'httpx',
|
|
65
|
+
'pydantic',
|
|
66
|
+
'tiktoken_ext',
|
|
67
|
+
'tiktoken_ext.openai_public',
|
|
68
|
+
],
|
|
69
|
+
hookspath=[],
|
|
70
|
+
hooksconfig={},
|
|
71
|
+
runtime_hooks=[],
|
|
72
|
+
excludes=[
|
|
73
|
+
'tkinter',
|
|
74
|
+
'matplotlib',
|
|
75
|
+
'scipy',
|
|
76
|
+
'numpy',
|
|
77
|
+
'pandas',
|
|
78
|
+
'pytest',
|
|
79
|
+
],
|
|
80
|
+
win_no_prefer_redirects=False,
|
|
81
|
+
win_private_assemblies=False,
|
|
82
|
+
cipher=block_cipher,
|
|
83
|
+
noarchive=False,
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
pyz = PYZ(a.pure, a.zipped_data, cipher=block_cipher)
|
|
87
|
+
|
|
88
|
+
exe = EXE(
|
|
89
|
+
pyz,
|
|
90
|
+
a.scripts,
|
|
91
|
+
a.binaries,
|
|
92
|
+
a.zipfiles,
|
|
93
|
+
a.datas,
|
|
94
|
+
[],
|
|
95
|
+
name='forge-core',
|
|
96
|
+
debug=False,
|
|
97
|
+
bootloader_ignore_signals=False,
|
|
98
|
+
strip=True,
|
|
99
|
+
upx=True,
|
|
100
|
+
upx_exclude=[],
|
|
101
|
+
runtime_tmpdir=None,
|
|
102
|
+
console=True,
|
|
103
|
+
disable_windowed_traceback=False,
|
|
104
|
+
argv_emulation=False,
|
|
105
|
+
target_arch=None,
|
|
106
|
+
codesign_identity=None,
|
|
107
|
+
entitlements_file=None,
|
|
108
|
+
)
|
|
File without changes
|
|
@@ -0,0 +1,87 @@
|
|
|
1
|
+
"""Load and template .md prompt files for AI consumption."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from pathlib import Path
|
|
6
|
+
|
|
7
|
+
from forge_core.utils import logger
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def load_prompt(prompts_dir: Path, prompt_name: str) -> str:
|
|
11
|
+
"""Load a prompt .md file from the prompts directory.
|
|
12
|
+
|
|
13
|
+
Args:
|
|
14
|
+
prompts_dir: Path to the prompts directory.
|
|
15
|
+
prompt_name: Name of the prompt file (e.g., 'detect-tech-stack').
|
|
16
|
+
|
|
17
|
+
Returns:
|
|
18
|
+
The prompt content as a string.
|
|
19
|
+
"""
|
|
20
|
+
file_path = prompts_dir / f"{prompt_name}.prompt.md"
|
|
21
|
+
if not file_path.exists():
|
|
22
|
+
logger.warn(f"Prompt not found: {file_path}")
|
|
23
|
+
return ""
|
|
24
|
+
|
|
25
|
+
content = file_path.read_text(encoding="utf-8")
|
|
26
|
+
# Strip YAML frontmatter if present
|
|
27
|
+
if content.startswith("---"):
|
|
28
|
+
parts = content.split("---", 2)
|
|
29
|
+
if len(parts) >= 3:
|
|
30
|
+
content = parts[2].strip()
|
|
31
|
+
|
|
32
|
+
return content
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def template_prompt(prompt: str, variables: dict[str, str]) -> str:
|
|
36
|
+
"""Replace template variables in a prompt.
|
|
37
|
+
|
|
38
|
+
Variables use {{variable_name}} syntax.
|
|
39
|
+
"""
|
|
40
|
+
result = prompt
|
|
41
|
+
for key, value in variables.items():
|
|
42
|
+
result = result.replace(f"{{{{{key}}}}}", value)
|
|
43
|
+
return result
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def build_file_context(files: dict[str, str], max_files: int = 50) -> str:
|
|
47
|
+
"""Build a file context string from a dict of {path: content}.
|
|
48
|
+
|
|
49
|
+
Formats each file with clear separators for the AI to parse.
|
|
50
|
+
"""
|
|
51
|
+
parts: list[str] = []
|
|
52
|
+
for i, (path, content) in enumerate(files.items()):
|
|
53
|
+
if i >= max_files:
|
|
54
|
+
parts.append(f"\n... and {len(files) - max_files} more files (truncated)")
|
|
55
|
+
break
|
|
56
|
+
parts.append(f"--- {path} ---\n{content}")
|
|
57
|
+
|
|
58
|
+
return "\n\n".join(parts)
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def load_learnings(central_path: str | None, local_path: Path | None) -> str:
|
|
62
|
+
"""Load LEARNINGS.md from central hub and/or local project."""
|
|
63
|
+
learnings_parts: list[str] = []
|
|
64
|
+
|
|
65
|
+
if central_path:
|
|
66
|
+
central_file = Path(central_path) / "LEARNINGS.md"
|
|
67
|
+
if central_file.exists():
|
|
68
|
+
learnings_parts.append(
|
|
69
|
+
f"=== Central Learnings ===\n{central_file.read_text(encoding='utf-8')}"
|
|
70
|
+
)
|
|
71
|
+
|
|
72
|
+
if local_path:
|
|
73
|
+
local_file = local_path / "LEARNINGS.md"
|
|
74
|
+
if local_file.exists():
|
|
75
|
+
learnings_parts.append(
|
|
76
|
+
f"=== Project Learnings ===\n{local_file.read_text(encoding='utf-8')}"
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
return "\n\n".join(learnings_parts)
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def load_knowledge_pack(central_path: str, pack_name: str) -> str:
|
|
83
|
+
"""Load a specific knowledge pack (e.g., kotlin-ktor.md)."""
|
|
84
|
+
pack_path = Path(central_path) / "knowledge-packs" / f"{pack_name}.md"
|
|
85
|
+
if pack_path.exists():
|
|
86
|
+
return pack_path.read_text(encoding="utf-8")
|
|
87
|
+
return ""
|
|
@@ -0,0 +1,121 @@
|
|
|
1
|
+
"""LiteLLM-based AI provider — unified interface for 100+ models."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
import litellm
|
|
8
|
+
|
|
9
|
+
from forge_core.models.config import AIConfig, AIProvider
|
|
10
|
+
from forge_core.utils import logger
|
|
11
|
+
from forge_core.utils.tokens import count_tokens
|
|
12
|
+
|
|
13
|
+
# Suppress LiteLLM's verbose logging
|
|
14
|
+
litellm.suppress_debug_info = True
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
def _resolve_model(config: AIConfig) -> str:
|
|
18
|
+
"""Resolve the model string for LiteLLM based on provider."""
|
|
19
|
+
if config.provider == AIProvider.OLLAMA:
|
|
20
|
+
return f"ollama/{config.model}"
|
|
21
|
+
if config.provider == AIProvider.AZURE:
|
|
22
|
+
return f"azure/{config.model}"
|
|
23
|
+
if config.provider == AIProvider.BEDROCK:
|
|
24
|
+
return f"bedrock/{config.model}"
|
|
25
|
+
return config.model
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def complete(
|
|
29
|
+
config: AIConfig,
|
|
30
|
+
system_prompt: str,
|
|
31
|
+
user_prompt: str,
|
|
32
|
+
json_mode: bool = False,
|
|
33
|
+
max_tokens: int | None = None,
|
|
34
|
+
) -> str:
|
|
35
|
+
"""Send a completion request to the configured AI provider.
|
|
36
|
+
|
|
37
|
+
Args:
|
|
38
|
+
config: AI provider configuration.
|
|
39
|
+
system_prompt: System-level instructions.
|
|
40
|
+
user_prompt: User message with code/context.
|
|
41
|
+
json_mode: If True, request JSON response format.
|
|
42
|
+
max_tokens: Override max tokens for this call.
|
|
43
|
+
|
|
44
|
+
Returns:
|
|
45
|
+
The AI's response text.
|
|
46
|
+
"""
|
|
47
|
+
model = _resolve_model(config)
|
|
48
|
+
messages = [
|
|
49
|
+
{"role": "system", "content": system_prompt},
|
|
50
|
+
{"role": "user", "content": user_prompt},
|
|
51
|
+
]
|
|
52
|
+
|
|
53
|
+
kwargs: dict[str, Any] = {
|
|
54
|
+
"model": model,
|
|
55
|
+
"messages": messages,
|
|
56
|
+
"temperature": config.temperature,
|
|
57
|
+
"max_tokens": max_tokens or config.max_tokens,
|
|
58
|
+
}
|
|
59
|
+
|
|
60
|
+
if config.api_key:
|
|
61
|
+
kwargs["api_key"] = config.api_key
|
|
62
|
+
if config.base_url:
|
|
63
|
+
kwargs["api_base"] = config.base_url
|
|
64
|
+
if json_mode:
|
|
65
|
+
kwargs["response_format"] = {"type": "json_object"}
|
|
66
|
+
|
|
67
|
+
# Log token usage
|
|
68
|
+
input_tokens = count_tokens(system_prompt + user_prompt, config.model)
|
|
69
|
+
logger.info(f"AI call → {model} ({input_tokens} input tokens)")
|
|
70
|
+
|
|
71
|
+
try:
|
|
72
|
+
response = litellm.completion(**kwargs)
|
|
73
|
+
content = response.choices[0].message.content or ""
|
|
74
|
+
|
|
75
|
+
output_tokens = count_tokens(content, config.model)
|
|
76
|
+
logger.info(f"AI response ← {output_tokens} output tokens")
|
|
77
|
+
|
|
78
|
+
return content
|
|
79
|
+
except Exception as e:
|
|
80
|
+
logger.error(f"AI call failed: {e}")
|
|
81
|
+
raise
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
def complete_with_fallback(
|
|
85
|
+
config: AIConfig,
|
|
86
|
+
system_prompt: str,
|
|
87
|
+
user_prompt: str,
|
|
88
|
+
fallback_models: list[str] | None = None,
|
|
89
|
+
json_mode: bool = False,
|
|
90
|
+
) -> str:
|
|
91
|
+
"""Try primary model, fall back to alternatives on failure.
|
|
92
|
+
|
|
93
|
+
Useful for complex classes where one model might fail.
|
|
94
|
+
"""
|
|
95
|
+
models = [_resolve_model(config)] + (fallback_models or [])
|
|
96
|
+
|
|
97
|
+
last_error = None
|
|
98
|
+
for model in models:
|
|
99
|
+
try:
|
|
100
|
+
kwargs: dict[str, Any] = {
|
|
101
|
+
"model": model,
|
|
102
|
+
"messages": [
|
|
103
|
+
{"role": "system", "content": system_prompt},
|
|
104
|
+
{"role": "user", "content": user_prompt},
|
|
105
|
+
],
|
|
106
|
+
"temperature": config.temperature,
|
|
107
|
+
"max_tokens": config.max_tokens,
|
|
108
|
+
}
|
|
109
|
+
if config.api_key:
|
|
110
|
+
kwargs["api_key"] = config.api_key
|
|
111
|
+
if json_mode:
|
|
112
|
+
kwargs["response_format"] = {"type": "json_object"}
|
|
113
|
+
|
|
114
|
+
response = litellm.completion(**kwargs)
|
|
115
|
+
return response.choices[0].message.content or ""
|
|
116
|
+
except Exception as e:
|
|
117
|
+
last_error = e
|
|
118
|
+
logger.warn(f"Model {model} failed, trying next: {e}")
|
|
119
|
+
continue
|
|
120
|
+
|
|
121
|
+
raise RuntimeError(f"All models failed. Last error: {last_error}")
|
|
@@ -0,0 +1,68 @@
|
|
|
1
|
+
"""Structured AI outputs using instructor + pydantic."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any, TypeVar
|
|
6
|
+
|
|
7
|
+
import instructor
|
|
8
|
+
import litellm
|
|
9
|
+
from pydantic import BaseModel
|
|
10
|
+
|
|
11
|
+
from forge_core.models.config import AIConfig
|
|
12
|
+
from forge_core.ai.provider import _resolve_model
|
|
13
|
+
from forge_core.utils import logger
|
|
14
|
+
|
|
15
|
+
T = TypeVar("T", bound=BaseModel)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def extract(
|
|
19
|
+
config: AIConfig,
|
|
20
|
+
system_prompt: str,
|
|
21
|
+
user_prompt: str,
|
|
22
|
+
response_model: type[T],
|
|
23
|
+
max_retries: int = 2,
|
|
24
|
+
) -> T:
|
|
25
|
+
"""Extract structured data from AI using instructor.
|
|
26
|
+
|
|
27
|
+
Forces the AI to return data matching a pydantic model schema.
|
|
28
|
+
Retries on validation failure.
|
|
29
|
+
|
|
30
|
+
Args:
|
|
31
|
+
config: AI provider configuration.
|
|
32
|
+
system_prompt: System-level instructions.
|
|
33
|
+
user_prompt: User message with code/context.
|
|
34
|
+
response_model: Pydantic model class to extract.
|
|
35
|
+
max_retries: Number of retries on validation failure.
|
|
36
|
+
|
|
37
|
+
Returns:
|
|
38
|
+
An instance of response_model populated by the AI.
|
|
39
|
+
"""
|
|
40
|
+
model = _resolve_model(config)
|
|
41
|
+
|
|
42
|
+
client = instructor.from_litellm(litellm.completion)
|
|
43
|
+
|
|
44
|
+
kwargs: dict[str, Any] = {
|
|
45
|
+
"model": model,
|
|
46
|
+
"messages": [
|
|
47
|
+
{"role": "system", "content": system_prompt},
|
|
48
|
+
{"role": "user", "content": user_prompt},
|
|
49
|
+
],
|
|
50
|
+
"temperature": config.temperature,
|
|
51
|
+
"max_retries": max_retries,
|
|
52
|
+
"response_model": response_model,
|
|
53
|
+
}
|
|
54
|
+
|
|
55
|
+
if config.api_key:
|
|
56
|
+
kwargs["api_key"] = config.api_key
|
|
57
|
+
if config.base_url:
|
|
58
|
+
kwargs["api_base"] = config.base_url
|
|
59
|
+
|
|
60
|
+
logger.info(f"Structured extraction → {model} → {response_model.__name__}")
|
|
61
|
+
|
|
62
|
+
try:
|
|
63
|
+
result = client.chat.completions.create(**kwargs)
|
|
64
|
+
logger.success(f"Extracted {response_model.__name__} successfully")
|
|
65
|
+
return result
|
|
66
|
+
except Exception as e:
|
|
67
|
+
logger.error(f"Structured extraction failed: {e}")
|
|
68
|
+
raise
|
|
@@ -0,0 +1,72 @@
|
|
|
1
|
+
"""SaaS API authentication and license verification."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import asyncio
|
|
6
|
+
from pathlib import Path
|
|
7
|
+
|
|
8
|
+
from forge_core.models.config import ForgeConfig, Plan, PlanLimits
|
|
9
|
+
from forge_core.utils import logger
|
|
10
|
+
from forge_core.utils.reporter import check_license
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
def verify_license(config: ForgeConfig) -> ForgeConfig:
|
|
14
|
+
"""Verify license with SaaS API and update plan limits.
|
|
15
|
+
|
|
16
|
+
If no auth token is set, runs in free/offline mode.
|
|
17
|
+
"""
|
|
18
|
+
if not config.auth_token:
|
|
19
|
+
logger.info("Running in offline mode (no auth token)")
|
|
20
|
+
config.limits = PlanLimits.for_plan(Plan.FREE)
|
|
21
|
+
return config
|
|
22
|
+
|
|
23
|
+
# Check with SaaS API
|
|
24
|
+
logger.info("Verifying license with SaaS API...")
|
|
25
|
+
try:
|
|
26
|
+
license_info = asyncio.run(check_license(config))
|
|
27
|
+
except Exception:
|
|
28
|
+
license_info = {}
|
|
29
|
+
|
|
30
|
+
if not license_info:
|
|
31
|
+
logger.warn("Could not verify license — falling back to free tier limits")
|
|
32
|
+
config.limits = PlanLimits.for_plan(Plan.FREE)
|
|
33
|
+
return config
|
|
34
|
+
|
|
35
|
+
# Update config from license response
|
|
36
|
+
plan_str = license_info.get("plan", "free")
|
|
37
|
+
try:
|
|
38
|
+
plan = Plan(plan_str)
|
|
39
|
+
except ValueError:
|
|
40
|
+
plan = Plan.FREE
|
|
41
|
+
|
|
42
|
+
config.limits = PlanLimits.for_plan(plan)
|
|
43
|
+
config.tenant.org_id = license_info.get("org_id", config.tenant.org_id)
|
|
44
|
+
config.tenant.org_name = license_info.get("org_name", config.tenant.org_name)
|
|
45
|
+
|
|
46
|
+
# If Pro/Enterprise, use SaaS proxy for AI calls
|
|
47
|
+
if plan in (Plan.PRO, Plan.ENTERPRISE):
|
|
48
|
+
config.ai.use_saas_proxy = True
|
|
49
|
+
if not config.ai.api_key:
|
|
50
|
+
config.ai.api_key = config.auth_token # SaaS proxy uses auth token
|
|
51
|
+
config.ai.base_url = f"{config.saas_api_url}/api/v1/ai"
|
|
52
|
+
|
|
53
|
+
logger.success(f"License verified: {plan.value} plan ({config.tenant.org_name})")
|
|
54
|
+
return config
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def save_auth_token(token: str) -> None:
|
|
58
|
+
"""Save auth token to ~/.forge-core/auth."""
|
|
59
|
+
auth_dir = Path.home() / ".forge-core"
|
|
60
|
+
auth_dir.mkdir(exist_ok=True)
|
|
61
|
+
auth_file = auth_dir / "auth"
|
|
62
|
+
auth_file.write_text(token, encoding="utf-8")
|
|
63
|
+
auth_file.chmod(0o600)
|
|
64
|
+
logger.success("Auth token saved")
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def load_auth_token() -> str:
|
|
68
|
+
"""Load auth token from ~/.forge-core/auth."""
|
|
69
|
+
auth_file = Path.home() / ".forge-core" / "auth"
|
|
70
|
+
if auth_file.exists():
|
|
71
|
+
return auth_file.read_text(encoding="utf-8").strip()
|
|
72
|
+
return ""
|