kiln-server 0.0.1__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.
- kiln_server-0.0.1/LICENSE.txt +4 -0
- kiln_server-0.0.1/PKG-INFO +44 -0
- kiln_server-0.0.1/README.md +22 -0
- kiln_server-0.0.1/kiln_server/custom_errors.py +70 -0
- kiln_server-0.0.1/kiln_server/project_api.py +122 -0
- kiln_server-0.0.1/kiln_server/prompt_api.py +36 -0
- kiln_server-0.0.1/kiln_server/provider_api.py +329 -0
- kiln_server-0.0.1/kiln_server/repair_api.py +83 -0
- kiln_server-0.0.1/kiln_server/run_api.py +156 -0
- kiln_server-0.0.1/kiln_server/server.py +57 -0
- kiln_server-0.0.1/kiln_server/settings.py +19 -0
- kiln_server-0.0.1/kiln_server/task_api.py +87 -0
- kiln_server-0.0.1/kiln_server/test_custom_error.py +94 -0
- kiln_server-0.0.1/kiln_server/test_project_api.py +483 -0
- kiln_server-0.0.1/kiln_server/test_prompt_api.py +86 -0
- kiln_server-0.0.1/kiln_server/test_provider_management.py +449 -0
- kiln_server-0.0.1/kiln_server/test_repair_api.py +182 -0
- kiln_server-0.0.1/kiln_server/test_run_api.py +715 -0
- kiln_server-0.0.1/kiln_server/test_server.py +208 -0
- kiln_server-0.0.1/kiln_server/test_settings.py +156 -0
- kiln_server-0.0.1/kiln_server/test_task_api.py +339 -0
- kiln_server-0.0.1/kiln_server/webhost.py +47 -0
- kiln_server-0.0.1/pyproject.toml +30 -0
|
@@ -0,0 +1,44 @@
|
|
|
1
|
+
Metadata-Version: 2.1
|
|
2
|
+
Name: kiln-server
|
|
3
|
+
Version: 0.0.1
|
|
4
|
+
Summary: Kiln AI Server
|
|
5
|
+
Home-page: https://kiln-ai.com
|
|
6
|
+
License: Proprietary
|
|
7
|
+
Author: Steve Cosman, Chesterfield Laboratories Inc
|
|
8
|
+
Author-email: scosman@users.noreply.github.com
|
|
9
|
+
Requires-Python: >=3.9
|
|
10
|
+
Classifier: License :: Other/Proprietary License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.9
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
15
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
16
|
+
Requires-Dist: kiln-ai
|
|
17
|
+
Project-URL: Bug Tracker, https://github.com/Kiln-AI/kiln/issues
|
|
18
|
+
Project-URL: Documentation, https://github.com/Kiln-AI/kiln#readme
|
|
19
|
+
Project-URL: Repository, https://github.com/Kiln-AI/kiln
|
|
20
|
+
Description-Content-Type: text/markdown
|
|
21
|
+
|
|
22
|
+
# Kiln AI Server
|
|
23
|
+
|
|
24
|
+
[](https://pypi.org/project/kiln-server)
|
|
25
|
+
[](https://pypi.org/project/kiln-server)
|
|
26
|
+
|
|
27
|
+
---
|
|
28
|
+
|
|
29
|
+
## About
|
|
30
|
+
|
|
31
|
+
The Kiln AI Server is a Python library that provides a REST API server for the Kiln AI datamodel.
|
|
32
|
+
|
|
33
|
+
See our [website](https://kiln-ai.com) for more information.
|
|
34
|
+
|
|
35
|
+
## Table of Contents
|
|
36
|
+
|
|
37
|
+
- [Installation](#installation)
|
|
38
|
+
|
|
39
|
+
## Installation
|
|
40
|
+
|
|
41
|
+
```console
|
|
42
|
+
pip install kiln_server
|
|
43
|
+
```
|
|
44
|
+
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
# Kiln AI Server
|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/kiln-server)
|
|
4
|
+
[](https://pypi.org/project/kiln-server)
|
|
5
|
+
|
|
6
|
+
---
|
|
7
|
+
|
|
8
|
+
## About
|
|
9
|
+
|
|
10
|
+
The Kiln AI Server is a Python library that provides a REST API server for the Kiln AI datamodel.
|
|
11
|
+
|
|
12
|
+
See our [website](https://kiln-ai.com) for more information.
|
|
13
|
+
|
|
14
|
+
## Table of Contents
|
|
15
|
+
|
|
16
|
+
- [Installation](#installation)
|
|
17
|
+
|
|
18
|
+
## Installation
|
|
19
|
+
|
|
20
|
+
```console
|
|
21
|
+
pip install kiln_server
|
|
22
|
+
```
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
from fastapi import FastAPI, HTTPException, Request, status
|
|
2
|
+
from fastapi.exceptions import RequestValidationError
|
|
3
|
+
from fastapi.responses import JSONResponse
|
|
4
|
+
from pydantic import ValidationError
|
|
5
|
+
|
|
6
|
+
|
|
7
|
+
def format_error_loc(loc: tuple):
|
|
8
|
+
if not loc:
|
|
9
|
+
return ""
|
|
10
|
+
formatted = []
|
|
11
|
+
# Skip the first item if it's "body" (case-insensitive)
|
|
12
|
+
start_index = 1 if loc and loc[0].lower() == "body" else 0
|
|
13
|
+
for i, item in enumerate(loc[start_index:]):
|
|
14
|
+
if item is not None and item != "":
|
|
15
|
+
if isinstance(item, str):
|
|
16
|
+
formatted.append(
|
|
17
|
+
item.capitalize() if i == 0 else "." + item.capitalize()
|
|
18
|
+
)
|
|
19
|
+
elif isinstance(item, int):
|
|
20
|
+
formatted.append(f"[{item}]")
|
|
21
|
+
else:
|
|
22
|
+
formatted.append(str(item))
|
|
23
|
+
return "".join(formatted)
|
|
24
|
+
|
|
25
|
+
|
|
26
|
+
def connect_custom_errors(app: FastAPI):
|
|
27
|
+
@app.exception_handler(ValidationError)
|
|
28
|
+
@app.exception_handler(RequestValidationError)
|
|
29
|
+
async def validation_exception_handler(
|
|
30
|
+
request: Request, exc: RequestValidationError | ValidationError
|
|
31
|
+
):
|
|
32
|
+
# Write user friendly error messages
|
|
33
|
+
error_messages = []
|
|
34
|
+
for error in exc.errors():
|
|
35
|
+
message = error.get("msg", "Unknown error")
|
|
36
|
+
loc = error.get("loc")
|
|
37
|
+
|
|
38
|
+
# Custom helpers for common errors
|
|
39
|
+
if "String should match pattern '^[A-Za-z0-9 _-]+$'" == message:
|
|
40
|
+
message = "must consist of only letters, numbers, spaces, hyphens, and underscores"
|
|
41
|
+
|
|
42
|
+
error_messages.append(f"{format_error_loc(loc)}: {message}")
|
|
43
|
+
|
|
44
|
+
def serialize_error(error):
|
|
45
|
+
return {
|
|
46
|
+
"type": error.get("type"),
|
|
47
|
+
"loc": [str(loc) for loc in error.get("loc", [])],
|
|
48
|
+
"msg": error.get("msg"),
|
|
49
|
+
"input": str(error.get("input")),
|
|
50
|
+
"ctx": {str(k): str(v) for k, v in error.get("ctx", {}).items()},
|
|
51
|
+
}
|
|
52
|
+
|
|
53
|
+
serialized_errors = [serialize_error(error) for error in exc.errors()]
|
|
54
|
+
|
|
55
|
+
return JSONResponse(
|
|
56
|
+
status_code=status.HTTP_422_UNPROCESSABLE_ENTITY,
|
|
57
|
+
content={
|
|
58
|
+
"error_messages": error_messages,
|
|
59
|
+
"message": ".\n".join(error_messages),
|
|
60
|
+
"source_errors": serialized_errors,
|
|
61
|
+
},
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
# Wrap in a format that the client can understand (message, and error_messages)
|
|
65
|
+
@app.exception_handler(HTTPException)
|
|
66
|
+
async def http_exception_handler(request: Request, exc: HTTPException):
|
|
67
|
+
return JSONResponse(
|
|
68
|
+
status_code=exc.status_code,
|
|
69
|
+
content={"message": exc.detail},
|
|
70
|
+
)
|
|
@@ -0,0 +1,122 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from pathlib import Path
|
|
3
|
+
from typing import Any, Dict
|
|
4
|
+
|
|
5
|
+
from fastapi import FastAPI, HTTPException
|
|
6
|
+
from kiln_ai.datamodel import Project
|
|
7
|
+
from kiln_ai.utils.config import Config
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def default_project_path():
|
|
11
|
+
return os.path.join(Path.home(), "Kiln Projects")
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def project_from_id(project_id: str) -> Project:
|
|
15
|
+
project_paths = Config.shared().projects
|
|
16
|
+
if project_paths is not None:
|
|
17
|
+
for project_path in project_paths:
|
|
18
|
+
try:
|
|
19
|
+
project = Project.load_from_file(project_path)
|
|
20
|
+
if project.id == project_id:
|
|
21
|
+
return project
|
|
22
|
+
except Exception as e:
|
|
23
|
+
print(f"Error loading project, skipping: {project_path}: {e}")
|
|
24
|
+
|
|
25
|
+
raise HTTPException(
|
|
26
|
+
status_code=404,
|
|
27
|
+
detail=f"Project not found. ID: {project_id}",
|
|
28
|
+
)
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def add_project_to_config(project_path: str):
|
|
32
|
+
projects = Config.shared().projects
|
|
33
|
+
if not isinstance(projects, list):
|
|
34
|
+
projects = []
|
|
35
|
+
if project_path not in projects:
|
|
36
|
+
projects.append(project_path)
|
|
37
|
+
Config.shared().save_setting("projects", projects)
|
|
38
|
+
|
|
39
|
+
|
|
40
|
+
def connect_project_api(app: FastAPI):
|
|
41
|
+
@app.post("/api/project")
|
|
42
|
+
async def create_project(project: Project) -> Project:
|
|
43
|
+
project_path = os.path.join(default_project_path(), project.name)
|
|
44
|
+
if os.path.exists(project_path):
|
|
45
|
+
raise HTTPException(
|
|
46
|
+
status_code=400,
|
|
47
|
+
detail="Project with this folder name already exists. Please choose a different name or rename the prior project's folder.",
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
os.makedirs(project_path)
|
|
51
|
+
project_file = os.path.join(project_path, "project.kiln")
|
|
52
|
+
project.path = Path(project_file)
|
|
53
|
+
project.save_to_file()
|
|
54
|
+
|
|
55
|
+
# add to projects list
|
|
56
|
+
add_project_to_config(project_file)
|
|
57
|
+
|
|
58
|
+
# Add path, which is usually excluded
|
|
59
|
+
return project
|
|
60
|
+
|
|
61
|
+
@app.patch("/api/project/{project_id}")
|
|
62
|
+
async def update_project(
|
|
63
|
+
project_id: str, project_updates: Dict[str, Any]
|
|
64
|
+
) -> Project:
|
|
65
|
+
original_project = project_from_id(project_id)
|
|
66
|
+
updated_project = original_project.model_copy(update=project_updates)
|
|
67
|
+
# Force validation using model_validate()
|
|
68
|
+
Project.model_validate(updated_project.model_dump())
|
|
69
|
+
updated_project.save_to_file()
|
|
70
|
+
return updated_project
|
|
71
|
+
|
|
72
|
+
@app.get("/api/projects")
|
|
73
|
+
async def get_projects() -> list[Project]:
|
|
74
|
+
project_paths = Config.shared().projects
|
|
75
|
+
projects = []
|
|
76
|
+
for project_path in project_paths if project_paths is not None else []:
|
|
77
|
+
try:
|
|
78
|
+
project = Project.load_from_file(project_path)
|
|
79
|
+
json_project = project.model_dump()
|
|
80
|
+
json_project["path"] = project_path
|
|
81
|
+
projects.append(json_project)
|
|
82
|
+
except Exception as e:
|
|
83
|
+
print(f"Error loading project, skipping: {project_path}: {e}")
|
|
84
|
+
|
|
85
|
+
return projects
|
|
86
|
+
|
|
87
|
+
@app.get("/api/projects/{project_id}")
|
|
88
|
+
async def get_project(project_id: str) -> Project:
|
|
89
|
+
return project_from_id(project_id)
|
|
90
|
+
|
|
91
|
+
# Removes the project, but does not delete the files from disk
|
|
92
|
+
@app.delete("/api/projects/{project_id}")
|
|
93
|
+
async def delete_project(project_id: str) -> dict:
|
|
94
|
+
project = project_from_id(project_id)
|
|
95
|
+
|
|
96
|
+
# Remove from config
|
|
97
|
+
projects_before = Config.shared().projects
|
|
98
|
+
projects_after = [p for p in projects_before if p != str(project.path)]
|
|
99
|
+
Config.shared().save_setting("projects", projects_after)
|
|
100
|
+
|
|
101
|
+
return {"message": f"Project removed. ID: {project_id}"}
|
|
102
|
+
|
|
103
|
+
@app.post("/api/import_project")
|
|
104
|
+
async def import_project(project_path: str) -> Project:
|
|
105
|
+
if project_path is None or not os.path.exists(project_path):
|
|
106
|
+
raise HTTPException(
|
|
107
|
+
status_code=400,
|
|
108
|
+
detail="Project not found. Check the path and try again.",
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
try:
|
|
112
|
+
project = Project.load_from_file(Path(project_path))
|
|
113
|
+
except Exception as e:
|
|
114
|
+
raise HTTPException(
|
|
115
|
+
status_code=500,
|
|
116
|
+
detail=f"Failed to load project. The file be invalid: {e}",
|
|
117
|
+
)
|
|
118
|
+
|
|
119
|
+
# add to projects list
|
|
120
|
+
add_project_to_config(project_path)
|
|
121
|
+
|
|
122
|
+
return project
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
from typing import Any, Dict, List
|
|
2
|
+
|
|
3
|
+
from fastapi import FastAPI, HTTPException
|
|
4
|
+
from kiln_ai.datamodel import Task
|
|
5
|
+
from pydantic import BaseModel
|
|
6
|
+
|
|
7
|
+
from libs.core.kiln_ai.adapters.prompt_builders import prompt_builder_from_ui_name
|
|
8
|
+
from libs.server.kiln_server.project_api import project_from_id
|
|
9
|
+
from libs.server.kiln_server.task_api import task_from_id
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
class PromptApiResponse(BaseModel):
|
|
13
|
+
prompt: str
|
|
14
|
+
prompt_builder_name: str
|
|
15
|
+
ui_generator_name: str
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def connect_prompt_api(app: FastAPI):
|
|
19
|
+
@app.get("/api/projects/{project_id}/task/{task_id}/gen_prompt/{prompt_generator}")
|
|
20
|
+
async def gen_prompt(
|
|
21
|
+
project_id: str, task_id: str, prompt_generator: str
|
|
22
|
+
) -> PromptApiResponse:
|
|
23
|
+
task = task_from_id(project_id, task_id)
|
|
24
|
+
|
|
25
|
+
try:
|
|
26
|
+
prompt_builder_class = prompt_builder_from_ui_name(prompt_generator)
|
|
27
|
+
prompt_builder = prompt_builder_class(task)
|
|
28
|
+
prompt = prompt_builder.build_prompt()
|
|
29
|
+
except Exception as e:
|
|
30
|
+
raise HTTPException(status_code=400, detail=str(e))
|
|
31
|
+
|
|
32
|
+
return PromptApiResponse(
|
|
33
|
+
prompt=prompt,
|
|
34
|
+
prompt_builder_name=prompt_builder_class.prompt_builder_name(),
|
|
35
|
+
ui_generator_name=prompt_generator,
|
|
36
|
+
)
|
|
@@ -0,0 +1,329 @@
|
|
|
1
|
+
import os
|
|
2
|
+
from typing import Dict, List
|
|
3
|
+
|
|
4
|
+
import requests
|
|
5
|
+
from fastapi import FastAPI, HTTPException
|
|
6
|
+
from fastapi.responses import JSONResponse
|
|
7
|
+
from kiln_ai.adapters.ml_model_list import (
|
|
8
|
+
ModelName,
|
|
9
|
+
ModelProviderName,
|
|
10
|
+
built_in_models,
|
|
11
|
+
provider_name_from_id,
|
|
12
|
+
provider_warnings,
|
|
13
|
+
)
|
|
14
|
+
from kiln_ai.utils.config import Config
|
|
15
|
+
from langchain_aws import ChatBedrockConverse
|
|
16
|
+
from pydantic import BaseModel
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class OllamaConnection(BaseModel):
|
|
20
|
+
message: str
|
|
21
|
+
models: List[str]
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
async def connect_ollama() -> OllamaConnection:
|
|
25
|
+
# Tags is a list of Ollama models. Proves Ollama is running, and models are available.
|
|
26
|
+
try:
|
|
27
|
+
tags = requests.get("http://localhost:11434/api/tags", timeout=5).json()
|
|
28
|
+
except requests.exceptions.ConnectionError:
|
|
29
|
+
raise HTTPException(
|
|
30
|
+
status_code=417,
|
|
31
|
+
detail="Failed to connect. Ensure Ollama app is running.",
|
|
32
|
+
)
|
|
33
|
+
except Exception as e:
|
|
34
|
+
raise HTTPException(
|
|
35
|
+
status_code=500,
|
|
36
|
+
detail=f"Failed to connect to Ollama: {e}",
|
|
37
|
+
)
|
|
38
|
+
|
|
39
|
+
# Build a list of models we support for Ollama from the built-in model list
|
|
40
|
+
supported_ollama_models = [
|
|
41
|
+
provider.provider_options["model"]
|
|
42
|
+
for model in built_in_models
|
|
43
|
+
for provider in model.providers
|
|
44
|
+
if provider.name == ModelProviderName.ollama
|
|
45
|
+
]
|
|
46
|
+
|
|
47
|
+
if "models" in tags:
|
|
48
|
+
models = tags["models"]
|
|
49
|
+
if isinstance(models, list):
|
|
50
|
+
model_names = [model["model"] for model in models]
|
|
51
|
+
available_supported_models = [
|
|
52
|
+
model
|
|
53
|
+
for model in model_names
|
|
54
|
+
if model in supported_ollama_models
|
|
55
|
+
or model in [f"{m}:latest" for m in supported_ollama_models]
|
|
56
|
+
]
|
|
57
|
+
if available_supported_models:
|
|
58
|
+
Config.shared().ollama_base_url = "http://localhost:11434"
|
|
59
|
+
return OllamaConnection(
|
|
60
|
+
message="Ollama connected",
|
|
61
|
+
models=available_supported_models,
|
|
62
|
+
)
|
|
63
|
+
|
|
64
|
+
raise HTTPException(
|
|
65
|
+
status_code=417,
|
|
66
|
+
detail="Ollama is running, but no supported models are installed. Install one or more supported model, like 'ollama pull phi3.5'.",
|
|
67
|
+
)
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
class ModelDetails(BaseModel):
|
|
71
|
+
id: str
|
|
72
|
+
name: str
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
class AvailableModels(BaseModel):
|
|
76
|
+
provider_name: str
|
|
77
|
+
provider_id: str
|
|
78
|
+
models: List[ModelDetails]
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
class ProviderModel(BaseModel):
|
|
82
|
+
id: str
|
|
83
|
+
name: str
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
class ProviderModels(BaseModel):
|
|
87
|
+
models: Dict[ModelName, ProviderModel]
|
|
88
|
+
|
|
89
|
+
|
|
90
|
+
def connect_provider_api(app: FastAPI):
|
|
91
|
+
@app.get("/api/providers/models")
|
|
92
|
+
async def get_providers_models() -> ProviderModels:
|
|
93
|
+
models = {}
|
|
94
|
+
for model in built_in_models:
|
|
95
|
+
models[model.name] = ProviderModel(id=model.name, name=model.friendly_name)
|
|
96
|
+
return ProviderModels(models=models)
|
|
97
|
+
|
|
98
|
+
# returns map, of provider name to list of model names
|
|
99
|
+
@app.get("/api/available_models")
|
|
100
|
+
async def get_available_models() -> List[AvailableModels]:
|
|
101
|
+
# Providers with just keys can return all their models if keys are set
|
|
102
|
+
key_providers: List[str] = []
|
|
103
|
+
|
|
104
|
+
# Try to connect to Ollama
|
|
105
|
+
try:
|
|
106
|
+
ollama_connection = await connect_ollama()
|
|
107
|
+
if len(ollama_connection.models) > 0:
|
|
108
|
+
key_providers.append(ModelProviderName.ollama)
|
|
109
|
+
except HTTPException:
|
|
110
|
+
# skip ollama if it's not available
|
|
111
|
+
pass
|
|
112
|
+
|
|
113
|
+
for provider, provider_warning in provider_warnings.items():
|
|
114
|
+
has_keys = True
|
|
115
|
+
for required_key in provider_warning.required_config_keys:
|
|
116
|
+
if Config.shared().get_value(required_key) is None:
|
|
117
|
+
has_keys = False
|
|
118
|
+
break
|
|
119
|
+
if has_keys:
|
|
120
|
+
key_providers.append(provider)
|
|
121
|
+
models: List[AvailableModels] = [
|
|
122
|
+
AvailableModels(
|
|
123
|
+
provider_name=provider_name_from_id(provider),
|
|
124
|
+
provider_id=provider,
|
|
125
|
+
models=[],
|
|
126
|
+
)
|
|
127
|
+
for provider in key_providers
|
|
128
|
+
]
|
|
129
|
+
|
|
130
|
+
for model in built_in_models:
|
|
131
|
+
for provider in model.providers:
|
|
132
|
+
if provider.name in key_providers:
|
|
133
|
+
available_models = next(
|
|
134
|
+
(m for m in models if m.provider_id == provider.name), None
|
|
135
|
+
)
|
|
136
|
+
if available_models:
|
|
137
|
+
available_models.models.append(
|
|
138
|
+
ModelDetails(id=model.name, name=model.friendly_name)
|
|
139
|
+
)
|
|
140
|
+
|
|
141
|
+
return models
|
|
142
|
+
|
|
143
|
+
@app.post("/api/provider/ollama/connect")
|
|
144
|
+
async def connect_ollama_api() -> OllamaConnection:
|
|
145
|
+
return await connect_ollama()
|
|
146
|
+
|
|
147
|
+
@app.post("/api/provider/connect_api_key")
|
|
148
|
+
async def connect_api_key(payload: dict):
|
|
149
|
+
provider = payload.get("provider")
|
|
150
|
+
key_data = payload.get("key_data")
|
|
151
|
+
if not isinstance(key_data, dict) or not isinstance(provider, str):
|
|
152
|
+
return JSONResponse(
|
|
153
|
+
status_code=400,
|
|
154
|
+
content={"message": "Invalid key_data or provider"},
|
|
155
|
+
)
|
|
156
|
+
|
|
157
|
+
api_key_providers = ["openai", "groq", "bedrock", "openrouter"]
|
|
158
|
+
if provider not in api_key_providers:
|
|
159
|
+
return JSONResponse(
|
|
160
|
+
status_code=400,
|
|
161
|
+
content={"message": f"Provider {provider} not supported"},
|
|
162
|
+
)
|
|
163
|
+
|
|
164
|
+
if provider == "openai" and isinstance(key_data["API Key"], str):
|
|
165
|
+
return await connect_openai(key_data["API Key"])
|
|
166
|
+
elif provider == "groq" and isinstance(key_data["API Key"], str):
|
|
167
|
+
return await connect_groq(key_data["API Key"])
|
|
168
|
+
elif provider == "openrouter" and isinstance(key_data["API Key"], str):
|
|
169
|
+
return await connect_openrouter(key_data["API Key"])
|
|
170
|
+
elif (
|
|
171
|
+
provider == "bedrock"
|
|
172
|
+
and isinstance(key_data["Access Key"], str)
|
|
173
|
+
and isinstance(key_data["Secret Key"], str)
|
|
174
|
+
):
|
|
175
|
+
return await connect_bedrock(key_data["Access Key"], key_data["Secret Key"])
|
|
176
|
+
else:
|
|
177
|
+
return JSONResponse(
|
|
178
|
+
status_code=400,
|
|
179
|
+
content={"message": f"Provider {provider} missing API key"},
|
|
180
|
+
)
|
|
181
|
+
|
|
182
|
+
|
|
183
|
+
async def connect_openrouter(key: str):
|
|
184
|
+
try:
|
|
185
|
+
headers = {
|
|
186
|
+
"Authorization": f"Bearer {key}",
|
|
187
|
+
"Content-Type": "application/json",
|
|
188
|
+
}
|
|
189
|
+
# invalid body, but we just want to see if the key is valid
|
|
190
|
+
response = requests.post(
|
|
191
|
+
"https://openrouter.ai/api/v1/chat/completions",
|
|
192
|
+
headers=headers,
|
|
193
|
+
json={},
|
|
194
|
+
)
|
|
195
|
+
|
|
196
|
+
# 401 def means invalid API key
|
|
197
|
+
if response.status_code == 401:
|
|
198
|
+
return JSONResponse(
|
|
199
|
+
status_code=401,
|
|
200
|
+
content={
|
|
201
|
+
"message": "Failed to connect to OpenRouter. Invalid API key."
|
|
202
|
+
},
|
|
203
|
+
)
|
|
204
|
+
else:
|
|
205
|
+
# No 401 means key is valid (even it it's an error, which we expect with empty body)
|
|
206
|
+
Config.shared().open_router_api_key = key
|
|
207
|
+
|
|
208
|
+
return JSONResponse(
|
|
209
|
+
status_code=200,
|
|
210
|
+
content={"message": "Connected to OpenRouter"},
|
|
211
|
+
)
|
|
212
|
+
# Any non-200 status code is an error
|
|
213
|
+
except Exception as e:
|
|
214
|
+
# unexpected error
|
|
215
|
+
return JSONResponse(
|
|
216
|
+
status_code=400,
|
|
217
|
+
content={"message": f"Failed to connect to OpenRouter. Error: {str(e)}"},
|
|
218
|
+
)
|
|
219
|
+
|
|
220
|
+
|
|
221
|
+
async def connect_openai(key: str):
|
|
222
|
+
try:
|
|
223
|
+
headers = {
|
|
224
|
+
"Authorization": f"Bearer {key}",
|
|
225
|
+
"Content-Type": "application/json",
|
|
226
|
+
}
|
|
227
|
+
response = requests.get("https://api.openai.com/v1/models", headers=headers)
|
|
228
|
+
|
|
229
|
+
# 401 def means invalid API key, so special case it
|
|
230
|
+
if response.status_code == 401:
|
|
231
|
+
return JSONResponse(
|
|
232
|
+
status_code=401,
|
|
233
|
+
content={"message": "Failed to connect to OpenAI. Invalid API key."},
|
|
234
|
+
)
|
|
235
|
+
|
|
236
|
+
# Any non-200 status code is an error
|
|
237
|
+
response.raise_for_status()
|
|
238
|
+
# If the request is successful, the function will continue
|
|
239
|
+
except Exception as e:
|
|
240
|
+
return JSONResponse(
|
|
241
|
+
status_code=400,
|
|
242
|
+
content={"message": f"Failed to connect to OpenAI. Error: {str(e)}"},
|
|
243
|
+
)
|
|
244
|
+
|
|
245
|
+
# It worked! Save the key and return success
|
|
246
|
+
Config.shared().open_ai_api_key = key
|
|
247
|
+
|
|
248
|
+
return JSONResponse(
|
|
249
|
+
status_code=200,
|
|
250
|
+
content={"message": "Connected to OpenAI"},
|
|
251
|
+
)
|
|
252
|
+
|
|
253
|
+
|
|
254
|
+
async def connect_groq(key: str):
|
|
255
|
+
try:
|
|
256
|
+
headers = {
|
|
257
|
+
"Authorization": f"Bearer {key}",
|
|
258
|
+
"Content-Type": "application/json",
|
|
259
|
+
}
|
|
260
|
+
response = requests.get(
|
|
261
|
+
"https://api.groq.com/openai/v1/models", headers=headers
|
|
262
|
+
)
|
|
263
|
+
|
|
264
|
+
if "invalid_api_key" in response.text:
|
|
265
|
+
return JSONResponse(
|
|
266
|
+
status_code=401,
|
|
267
|
+
content={"message": "Failed to connect to Groq. Invalid API key."},
|
|
268
|
+
)
|
|
269
|
+
|
|
270
|
+
# Any non-200 status code is an error
|
|
271
|
+
response.raise_for_status()
|
|
272
|
+
# If the request is successful, the function will continue
|
|
273
|
+
except Exception as e:
|
|
274
|
+
return JSONResponse(
|
|
275
|
+
status_code=400,
|
|
276
|
+
content={"message": f"Failed to connect to Groq. Error: {str(e)}"},
|
|
277
|
+
)
|
|
278
|
+
|
|
279
|
+
# It worked! Save the key and return success
|
|
280
|
+
Config.shared().groq_api_key = key
|
|
281
|
+
|
|
282
|
+
return JSONResponse(
|
|
283
|
+
status_code=200,
|
|
284
|
+
content={"message": "Connected to Groq"},
|
|
285
|
+
)
|
|
286
|
+
|
|
287
|
+
|
|
288
|
+
async def connect_bedrock(access_key: str, secret_key: str):
|
|
289
|
+
try:
|
|
290
|
+
# Langchain API is not good... need to use env vars. Pop these in finally block
|
|
291
|
+
os.environ["AWS_ACCESS_KEY_ID"] = access_key
|
|
292
|
+
os.environ["AWS_SECRET_ACCESS_KEY"] = secret_key
|
|
293
|
+
|
|
294
|
+
# Fake model, but will get a credential error before AWS realizes it's wrong
|
|
295
|
+
# Ensures we don't spend money on a test call
|
|
296
|
+
llm = ChatBedrockConverse(
|
|
297
|
+
model="fake_model",
|
|
298
|
+
region_name="us-west-2",
|
|
299
|
+
)
|
|
300
|
+
llm.invoke("Hello, how are you?")
|
|
301
|
+
except Exception as e:
|
|
302
|
+
# Check for specific error messages indicating invalid credentials
|
|
303
|
+
if "UnrecognizedClientException" in str(
|
|
304
|
+
e
|
|
305
|
+
) or "InvalidSignatureException" in str(e):
|
|
306
|
+
return JSONResponse(
|
|
307
|
+
status_code=401,
|
|
308
|
+
content={
|
|
309
|
+
"message": "Failed to connect to Bedrock. Invalid credentials."
|
|
310
|
+
},
|
|
311
|
+
)
|
|
312
|
+
else:
|
|
313
|
+
# We still expect an error (fake model), but not for invalid credentials which means success!
|
|
314
|
+
Config.shared().bedrock_access_key = access_key
|
|
315
|
+
Config.shared().bedrock_secret_key = secret_key
|
|
316
|
+
return JSONResponse(
|
|
317
|
+
status_code=200,
|
|
318
|
+
content={"message": "Connected to Bedrock"},
|
|
319
|
+
)
|
|
320
|
+
|
|
321
|
+
finally:
|
|
322
|
+
os.environ.pop("AWS_ACCESS_KEY_ID")
|
|
323
|
+
os.environ.pop("AWS_SECRET_ACCESS_KEY")
|
|
324
|
+
|
|
325
|
+
# we shouldn't get here, but if we do, something went wrong
|
|
326
|
+
return JSONResponse(
|
|
327
|
+
status_code=400,
|
|
328
|
+
content={"message": "Unknown Bedrock Error"},
|
|
329
|
+
)
|