diff --git a/.gitea/act-runner/.env.example b/.gitea/act-runner/.env.example index 6717d4b..edf20d9 100644 --- a/.gitea/act-runner/.env.example +++ b/.gitea/act-runner/.env.example @@ -1,3 +1,3 @@ GITEA_INSTANCE_URL=https://git.example.com GITEA_RUNNER_REGISTRATION_TOKEN=从仓库_Settings_Actions_Runners_复制 -GITEA_RUNNER_NAME=go-vue-template-ci-runner +GITEA_RUNNER_NAME=atlas-ci-runner diff --git a/.gitea/act-runner/docker-compose.yml b/.gitea/act-runner/docker-compose.yml index c652b98..c5b8ecd 100644 --- a/.gitea/act-runner/docker-compose.yml +++ b/.gitea/act-runner/docker-compose.yml @@ -5,13 +5,13 @@ services: act-runner: image: docker.io/gitea/act_runner:0.2.12 - container_name: go-vue-template-act-runner + container_name: atlas-act-runner restart: unless-stopped environment: CONFIG_FILE: /config.yaml GITEA_INSTANCE_URL: ${GITEA_INSTANCE_URL:-https://git.example.com} GITEA_RUNNER_REGISTRATION_TOKEN: ${GITEA_RUNNER_REGISTRATION_TOKEN} - GITEA_RUNNER_NAME: ${GITEA_RUNNER_NAME:-go-vue-template-ci-runner} + GITEA_RUNNER_NAME: ${GITEA_RUNNER_NAME:-atlas-ci-runner} volumes: - ./config.yaml:/config.yaml:ro - act-runner-data:/data diff --git a/Dockerfile b/Dockerfile index d60cfbb..d516a84 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,5 +1,3 @@ -# syntax=docker/dockerfile:1 - ARG IMAGE_PREFIX= FROM ${IMAGE_PREFIX}node:20-alpine AS web-builder diff --git a/Makefile b/Makefile index b7b12f7..1fa3ecc 100644 --- a/Makefile +++ b/Makefile @@ -3,7 +3,7 @@ export GOPROXY ?= https://goproxy.cn,direct export GOSUMDB ?= sum.golang.google.cn export IMAGE_PREFIX ?= docker.1panel.live/library/ -export APP_IMAGE ?= go-vue-template:latest +export APP_IMAGE ?= atlas:latest dev: @echo "Run backend and frontend in separate terminals:" @@ -32,7 +32,7 @@ test: clean: rm -rf dist web/dist internal/api/static/* - @echo 'App' > internal/api/static/index.html + @echo 'AtlasAtlas' > internal/api/static/index.html docker-build: docker build \ diff --git a/README.md b/README.md index 4bd56cb..6eb64ac 100644 --- a/README.md +++ b/README.md @@ -1,14 +1,16 @@ -# go-vue-template +# Atlas -Go + Vue 单体应用脚手架。一个二进制同时提供 REST API 与管理台 UI，前端通过 `go:embed` 嵌入，适合内部工具、管理后台、网关控制台等场景。 +Atlas 是面向 AI 能力的统一接入与使用平台。Go + Vue 单体部署：用户端提供对话、文生图、图生图与历史记录；管理端配置上游 Provider 与模型、查看请求日志；内部 Gateway 兼容 OpenAI / Replicate 协议，统一编排、日志与路由。 ## 特性 - **单体部署**：API 与 SPA 同进程、同端口，生产环境无需 Node.js -- **分层清晰**：`cmd` → `api` → `service` → `store`，职责边界明确 -- **开箱即用**：登录鉴权、设置页、审计日志、i18n、Docker 多阶段构建 -- **纯 Go 编译**：SQLite 使用 `modernc.org/sqlite`，`CGO_ENABLED=0` 可交叉编译 -- **可演进**：goose SQL 迁移、settings KV 配置、Pinia 状态管理 +- **多协议 Gateway**：`/v1/*` 透传 OpenAI（Chat、Embeddings 等）与 Replicate（Predictions） +- **统一 invoke 编排**：同步/异步调用、任务队列、上游轮询 + 前端 SSE 进度推送 +- **Provider 管理**：管理端配置上游地址与密钥；模型优先从上游 `/v1/models` 同步，也可手动维护 +- **请求日志**：Gateway 与用户 invoke 全链路记录，管理端可检索 +- **用户会话隔离**：Cookie 会话区分用户数据；管理端独立 Bearer 鉴权 +- **纯 Go 编译**：SQLite（modernc.org/sqlite），`CGO_ENABLED=0` 可交叉编译 ## 快速开始 @@ -25,63 +27,79 @@ make build docker compose up -d ``` -默认管理员密码：`admin`（可在设置页修改）。 +启动后访问： ---- +| 入口 | 地址 | 说明 | +|------|------|------| +| 用户端 | `http://localhost:8080/chat` | 对话、生图、历史（无需登录） | +| 管理端 | `http://localhost:8080/admin` | Provider、日志、设置（需登录） | +| Gateway | `http://localhost:8080/v1/*` | OpenAI / Replicate 兼容 API | +| 健康检查 | `GET /health` | `{ "status": "ok" }` | -## 使用本模板 - -1. **Fork 或 clone** 本仓库作为新项目起点 -2. **修改 Go module 名**：全局替换 `github.com/rose_cat707/go-vue-template` 为你的模块路径 -3. **修改应用标识**：更新 `settings` 默认值、`web/` 中的品牌文案、镜像名等 -4. **按需扩展**：在 `internal/` 增加业务包，在 `web/src/views/` 增加页面 - -```bash -# 示例：批量替换 module 名 -find . -type f \( -name '*.go' -o -name 'go.mod' \) \ - -exec sed -i '' 's|github.com/rose_cat707/go-vue-template|github.com/you/your-app|g' {} + -``` +默认管理员密码：`admin`（可在管理端设置页修改）。 --- ## 架构 ``` -┌─────────────────────────────────────────────────────────┐ -│ 单一 Go 进程 │ -│ ┌──────────────┐ ┌──────────────┐ ┌───────────────┐ │ -│ │ /api/v1/* │ │ /health │ │ /* (SPA) │ │ -│ │ REST API │ │ 健康检查 │ │ embed 静态资源 │ │ -│ └──────┬───────┘ └──────────────┘ └───────────────┘ │ -│ │ │ -│ ┌──────▼──────────────────────────────────────────┐ │ -│ │ api (HTTP) → service (编排) → store (持久化) │ │ -│ └─────────────────────────────────────────────────┘ │ -│ │ │ -│ ┌──────▼──────┐ │ -│ │ SQLite WAL │ {data}/app.db │ -│ └─────────────┘ │ -└─────────────────────────────────────────────────────────┘ - ▲ ▲ - │ 生产：同一端口 │ 开发：Vite :5173 代理 /api - └──────────────────────────┘ +┌─────────────────────────────────────────────────────────────────────┐ +│ 单一 Go 进程 │ +│ ┌─────────────┐ ┌─────────────┐ ┌──────────────┐ ┌──────────┐ │ +│ │ /api/v1/* │ │ /v1/* │ │ /health │ │ /* SPA │ │ +│ │ 用户 + 管理 │ │ Gateway │ │ │ │ embed │ │ +│ └──────┬──────┘ └──────┬──────┘ └──────────────┘ └──────────┘ │ +│ │ │ │ +│ ┌──────▼────────────────▼──────────────────────────────────────┐ │ +│ │ api → service (invoke / chat / reqlog / taskevents) → store │ │ +│ └──────────────────────────┬───────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────▼───────────────────────────────────┐ │ +│ │ provider (Router / ModelsService / ProxyStream) → 上游 API │ │ +│ └────────────────────────────────────────────────────────────────┘ │ +│ │ │ +│ ┌──────────────────────────▼───────────────────────────────────┐ │ +│ │ SQLite WAL {data}/app.db │ │ +│ └────────────────────────────────────────────────────────────────┘ │ +└─────────────────────────────────────────────────────────────────────┘ +``` + +### 调用链路 + +**对话（Chat）** + +``` +浏览器 → POST /api/v1/user/conversations/:id/messages + → chat.Service → provider.ProxyStream → 上游 SSE 真流式回传 +``` + +**生图（异步任务）** + +``` +浏览器 → POST /api/v1/user/generate + → invoke.Service 创建 task → async_worker 轮询上游 + → GET /api/v1/user/tasks/:id/events (SSE) 推送进度 +``` + +**Gateway 直连** + +``` +客户端 → /v1/chat/completions | /v1/predictions ... + → provider.Router 按模型路由 → 上游 + request_logs ``` ### 分层约定 | 层 | 路径 | 职责 | |----|------|------| -| 入口 | `cmd/app/` | 解析参数、初始化依赖、启动 HTTP、优雅关闭 | -| HTTP | `internal/api/` | 路由注册、中间件、请求/响应、SPA 静态服务 | -| 业务 | `internal/service/` | 跨模块编排，注入 Store 与领域服务 | -| 持久化 | `internal/store/` | 数据库连接、迁移、仓储方法 | -| 配置 | `internal/settings/` | 配置键常量、默认值、展示脱敏、提交校验 | -| 前端 | `web/` | Vue SPA 源码，构建产物复制到 `internal/api/static/` | - -**原则：** -- Handler 只做 HTTP 适配，不写复杂业务逻辑 -- Store 封装 SQL，不暴露 `*sql.DB` 给 Handler -- 前端构建产物不提交 dist，仅保留 `static/index.html` 占位以保证 `go test` / CI 可通过 +| 入口 | `cmd/app/` | 参数解析、依赖初始化、HTTP 服务、优雅关闭 | +| HTTP | `internal/api/` | 路由、中间件、用户/管理/Gateway Handler、SPA | +| 编排 | `internal/service/` | invoke、chat、reqlog、taskevents、asset | +| 上游 | `internal/provider/` | HTTP 代理、模型列表、流式转发、Provider 路由 | +| 领域 | `internal/domain/` | 服务类型、invoke 请求/响应模型 | +| 持久化 | `internal/store/` | SQLite、goose 迁移、仓储与 HotCache | +| 配置 | `internal/settings/` | KV 配置键、默认值、脱敏 | +| 前端 | `web/` | Vue SPA，构建产物嵌入 `internal/api/static/` | --- @@ -89,96 +107,122 @@ find . -type f \( -name '*.go' -o -name 'go.mod' \) \ ``` . -├── cmd/app/ # 进程入口 +├── cmd/app/ # 进程入口 ├── internal/ │ ├── api/ -│ │ ├── handlers.go # 路由与 Handler -│ │ ├── auth_token.go # Token 签发与校验 -│ │ ├── static.go # go:embed SPA -│ │ └── static/ # 前端构建产物（占位 index.html 已提交） -│ ├── service/ # 业务编排 -│ ├── store/ -│ │ ├── store.go -│ │ └── migrations/ # goose SQL 迁移 -│ └── settings/ # 配置键与规范化 -├── web/ -│ ├── src/ -│ │ ├── api/ # axios 客户端与类型 -│ │ ├── components/ # 通用 UI 组件 -│ │ ├── i18n/ # 多语言 -│ │ ├── layouts/ # 布局 -│ │ ├── router/ # 路由与鉴权守卫 -│ │ ├── stores/ # Pinia -│ │ ├── styles/ # 主题变量 -│ │ └── views/ # 页面 -│ ├── vite.config.ts -│ └── package.json +│ │ ├── handlers.go # 基础路由、鉴权、设置 +│ │ ├── user_handlers.go # 用户端 API +│ │ ├── admin_handlers.go # 管理端 API +│ │ ├── gateway*.go # /v1 Gateway +│ │ ├── task_sse.go # 任务 SSE +│ │ └── static/ # 前端 embed 产物 +│ ├── service/ +│ │ ├── invoke/ # 同步/异步 invoke + worker +│ │ ├── chat/ # 对话编排 +│ │ ├── reqlog/ # 请求日志 +│ │ └── taskevents/ # 任务事件广播 +│ ├── provider/ # 上游 HTTP 代理与路由 +│ ├── domain/ +│ ├── store/migrations/ # goose SQL（002_atlas, 003_provider_models） +│ └── settings/ +├── web/src/ +│ ├── views/ # Chat、生图、History、Admin 页面 +│ ├── layouts/ # UserLayout / AdminLayout +│ ├── composables/ # useInvoke (SSE)、useUserMeta 等 +│ └── styles/global.css # 设计规范落地样式 +├── docs/UI-DESIGN-SYSTEM.md # UI 设计规范 ├── Makefile ├── Dockerfile -├── docker-compose.yml -├── go.mod -└── .gitea/ # Gitea Actions CI - ├── README.md - ├── workflows/ci.yml - └── act-runner/ # runner 部署参考（可选） +└── docker-compose.yml ``` --- ## 技术栈 -| 层级 | 选型 | 说明 | -|------|------|------| -| 后端 | Go 1.25 + Gin | HTTP 路由与中间件 | -| 数据库 | SQLite (modernc.org/sqlite) | 纯 Go 驱动，无 CGO | -| 迁移 | goose v3 (embed SQL) | 版本化、可回滚 | -| 前端 | Vue 3 + Vite + TypeScript | Composition API + ` + + +``` + +### 8.2 登录页完整模板 + +```vue + + + + +``` + +### 8.3 公共组件完整源码 + +```vue + + + + +``` + +```vue + + + + +``` + +```vue + + + + +``` + +品牌 Logo 由各项目自行提供，放入 `.brand-icon`（20px）与 `.login-brand-icon`（24px）容器即可；须使用 `currentColor` 或适配主色底上的 `--primary-foreground`。 + +### 8.4 标准页面模板 + +```vue + +``` + +### 8.5 带 Tab 的列表页 + +```vue + +
+ + + + +
+
…
+ … +
…
+ +``` + +### 8.6 新页面 Checklist + +1. 将 §8.9 复制为项目全局样式文件，按 §6.0 顺序在 `main.ts` 中引入 +2. 复制 §8.1 壳层布局、§8.3 公共组件（文件名自定） +3. 页面根节点 `
` +4. `PageHeader`（标题 + 描述 + actions） +5. 主内容 `AppCard`；表单页加 `padded` +6. 列表筛选用 `.card-toolbar`，分页用 `.card-pager` +7. 间距与颜色仅使用 CSS 变量与全局类 +8. 在主布局中注册导航项与图标（§5.3） + +### 8.7 反模式 + +| 不要 | 应该 | +|------|------| +| 深色侧栏 | 浅色半透明侧栏 | +| 页面内裸 `el-card` 作主容器 | `AppCard` | +| 无标题直接上表格 | 先 `PageHeader` | +| 硬编码色值 | CSS 变量 | +| 重阴影 | 细边框 | +| 每页自定义 `.toolbar` / `.pager` | `.card-toolbar` / `.card-pager` | +| 表单卡片不加 `padded` | `AppCard padded` | +| `style="margin-top: 12px"` | `.u-mt-3` 等工具类 | +| 非 8px 倍数的随意间距 | 间距 Token | +| 在 EP 样式之前引入全局样式 | 先 `element-plus/dist/index.css`，再引入 §8.9 全局样式 | + +### 8.8 推荐目录结构（示例） + +``` +styles/global.css ← 由本文档 §8.9 复制生成，路径与文件名自定 +components/PageHeader.vue +components/StatCard.vue +components/AppCard.vue +layouts/MainLayout.vue ← 见 §8.1，文件名自定 +views/Login.vue ← 见 §8.2，文件名自定 +``` + +样式变更时，须同时更新本文档 **§6.1–6.14** 与 **§8.9** 两处 CSS，保持内容一致。 + +### 8.9 完整全局样式（一键复制） + +以下为**设计规范定义的完整全局 CSS**，直接复制到项目全局样式文件即可（§6.1–6.14 为分模块对照版，内容须与本节一致）。 + +```css +:root { + --background: #f4f4f5; + --foreground: #18181b; + --card: #ffffff; + --card-foreground: #18181b; + --primary: #0f7a62; + --primary-hover: #0d6b56; + --primary-foreground: #fafafa; + --muted: #f4f4f5; + --muted-foreground: #71717a; + --border: #e4e4e7; + --input: #e4e4e7; + --destructive: #dc2626; + --success: #047857; + --warning: #b45309; + --danger: #dc2626; + --info: #a1a1aa; + --sidebar: rgba(250, 250, 250, 0.85); + --sidebar-foreground: #18181b; + --sidebar-accent: #f4f4f5; + --sidebar-border: #e4e4e7; + --sidebar-width: 15rem; + --radius: 0.5rem; + --space-1: 4px; + --space-2: 8px; + --space-3: 12px; + --space-4: 16px; + --space-5: 20px; + --space-6: 24px; + --space-8: 32px; + + --el-color-primary: #0f7a62; + --el-color-primary-light-3: #4da08c; + --el-color-primary-light-5: #76b8a8; + --el-color-primary-light-7: #a0d0c4; + --el-color-primary-light-8: #b8dcd2; + --el-color-primary-light-9: #d0e8e1; + --el-color-primary-dark-2: #0c624e; + --el-color-success: #047857; + --el-color-warning: #b45309; + --el-color-danger: #dc2626; + --el-color-info: #a1a1aa; + --el-border-radius-base: 0.5rem; + --el-border-color: #e4e4e7; + --el-bg-color-page: #f4f4f5; + --el-text-color-primary: #18181b; + --el-text-color-regular: #3f3f46; + --el-text-color-secondary: #71717a; + --el-font-family: ui-sans-serif, system-ui, -apple-system, "Segoe UI", Roboto, "Helvetica Neue", Arial, "PingFang SC", "Microsoft YaHei", "Noto Sans SC", sans-serif; +} + +* { box-sizing: border-box; } + +html, body, #app { + margin: 0; + padding: 0; + height: 100%; + font-size: 14px; + font-family: var(--el-font-family); + background: var(--background); + color: var(--foreground); +} + +.layout { + display: flex; + height: 100vh; +} + +.sidebar { + width: var(--sidebar-width); + flex-shrink: 0; + background: var(--sidebar); + border-right: 1px solid var(--sidebar-border); + backdrop-filter: blur(8px); + display: flex; + flex-direction: column; + overflow-y: auto; +} + +.sidebar-brand { + display: flex; + align-items: center; + gap: var(--space-3); + padding: var(--space-5) var(--space-4); +} + +.brand-icon { + width: 32px; + height: 32px; + border-radius: var(--radius); + background: var(--primary); + color: var(--primary-foreground); + display: flex; + align-items: center; + justify-content: center; +} + +.brand-name { + font-size: 1.125rem; + font-weight: 600; +} + +.nav-group-label { + font-size: 0.75rem; + text-transform: uppercase; + letter-spacing: 0.04em; + color: var(--muted-foreground); + padding: var(--space-4) var(--space-4) var(--space-2); +} + +.nav-item { + display: flex; + align-items: center; + gap: var(--space-2); + padding: var(--space-2) var(--space-3); + margin: 0 var(--space-2); + border-radius: var(--radius); + color: var(--muted-foreground); + text-decoration: none; + font-size: 0.875rem; + font-weight: 500; +} + +.nav-item:hover { background: var(--sidebar-accent); } + +.nav-item.active { + background: var(--sidebar-accent); + color: var(--primary); + font-weight: 600; +} + +.main-area { + flex: 1; + display: flex; + flex-direction: column; + min-width: 0; +} + +.topbar { + height: 52px; + display: flex; + align-items: center; + justify-content: space-between; + padding: 0 var(--space-6); + background: var(--card); + border-bottom: 1px solid var(--border); +} + +.topbar-title { + font-size: 0.875rem; + color: var(--muted-foreground); +} + +.main-content { + flex: 1; + overflow-y: auto; + padding: var(--space-6); +} + +.page { + max-width: 1400px; +} + +.page-header { + display: flex; + align-items: flex-start; + justify-content: space-between; + margin-bottom: var(--space-6); + gap: var(--space-4); +} + +.page-actions { + display: flex; + align-items: center; + gap: var(--space-2); + flex-shrink: 0; + flex-wrap: wrap; +} + +.page-title { + font-size: 1.5rem; + font-weight: 600; + letter-spacing: -0.02em; + margin: 0 0 var(--space-1); +} + +.page-desc { + font-size: 0.875rem; + color: var(--muted-foreground); + margin: 0; +} + +.stat-grid { + display: grid; + grid-template-columns: repeat(4, 1fr); + gap: var(--space-4); + margin-bottom: var(--space-6); +} + +.stat-card { + background: var(--card); + border: 1px solid var(--border); + border-radius: var(--radius); + padding: var(--space-5); +} + +.stat-card__label { + font-size: 0.875rem; + color: var(--muted-foreground); + margin-bottom: var(--space-2); +} + +.stat-card__value { + font-size: 1.75rem; + font-weight: 600; +} + +.app-card { + background: var(--card); + border: 1px solid var(--border); + border-radius: var(--radius); + margin-bottom: var(--space-4); +} + +.page > .app-card:last-child, +.page > .content-grid:last-child { + margin-bottom: 0; +} + +.app-card__header { + padding: var(--space-4) var(--space-5); + border-bottom: 1px solid var(--border); +} + +.app-card__title { + font-size: 0.9375rem; + font-weight: 600; + margin: 0; +} + +.app-card__body--padded { padding: var(--space-5); } + +.content-grid { + display: grid; + gap: var(--space-4); + margin-bottom: var(--space-6); +} + +.content-grid--2-1 { grid-template-columns: 1.4fr 1fr; } +.content-grid--2 { grid-template-columns: 1fr 1fr; } +.content-grid--3 { grid-template-columns: repeat(3, 1fr); } + +.card-toolbar { + display: flex; + align-items: center; + gap: var(--space-2); + padding: var(--space-4) var(--space-5) 0; + margin-bottom: var(--space-4); + flex-wrap: wrap; +} + +.card-toolbar--desc { + align-items: flex-start; +} + +.card-pager { + padding: var(--space-4) var(--space-5); + display: flex; + justify-content: flex-end; +} + +.card-footer-hint { + padding: 0 var(--space-5) var(--space-3); + margin: 0; +} + +.card-footer-hint + .card-pager { + padding-top: var(--space-3); +} + +.card-inset-block { + margin: 0 var(--space-5) var(--space-4); + padding: var(--space-3); + background: var(--muted); + border-radius: var(--radius); + font-size: 0.875rem; +} + +.tabs-bar { + padding: 0 var(--space-5); + border-bottom: 1px solid var(--border); +} + +.tabs-bar .el-tabs__header { margin-bottom: 0; } + +.page-toolbar { + display: flex; + align-items: center; + gap: var(--space-2); + margin-bottom: var(--space-4); + flex-wrap: wrap; +} + +.inline-actions { + display: flex; + align-items: center; + gap: var(--space-2); + flex-wrap: wrap; +} + +.card-header-row { + display: flex; + align-items: center; + justify-content: space-between; + gap: var(--space-4); +} + +.form-hint { + margin-top: 6px; + font-size: 0.75rem; + color: var(--muted-foreground); +} + +.form-hint--flush { margin-top: 0; } + +.u-mt-3 { margin-top: var(--space-3); } +.u-mt-4 { margin-top: var(--space-4); } +.u-mb-3 { margin-bottom: var(--space-3); } +.u-mb-4 { margin-bottom: var(--space-4); } +.u-mb-6 { margin-bottom: var(--space-6); } +.text-center { text-align: center; } +.w-full { width: 100%; } + +code.inline-code { + font-size: 0.75rem; + padding: 0 var(--space-1); + background: var(--muted); + border-radius: var(--radius); +} + +.section-title { + font-size: 0.875rem; + font-weight: 600; + margin-bottom: var(--space-4); +} + +.section-title:not(:first-child) { margin-top: var(--space-6); } + +.login-page { + min-height: 100vh; + display: flex; + align-items: center; + justify-content: center; + background: var(--background); + padding: var(--space-6); +} + +.login-wrap { max-width: 400px; width: 100%; text-align: center; } + +.login-brand-icon { + width: 48px; + height: 48px; + border-radius: 12px; + background: var(--primary); + color: var(--primary-foreground); + display: inline-flex; + align-items: center; + justify-content: center; + margin-bottom: var(--space-4); +} + +.login-card { + background: var(--card); + border: 1px solid var(--border); + border-radius: var(--radius); + padding: var(--space-8); + margin-top: var(--space-6); + text-align: left; +} + +.status-good { color: #047857; } +.status-warn { color: #b45309; } +.status-bad { color: #dc2626; } +.status-muted { color: var(--muted-foreground); } + +.toolbar-select { width: 120px; } +.toolbar-select--country { width: 160px; } +.toolbar-select--wide { width: 280px; } +.toolbar-input { width: 200px; } +.toolbar-input--wide { width: 260px; } +.toolbar-input--grow { flex: 1; min-width: 200px; } +.table-cell-select { width: 100%; min-width: 200px; } + +.el-form .el-select { width: 100%; } + +.app-card .el-table { + --el-table-header-bg-color: var(--muted); + font-size: 0.875rem; +} + +.app-card .el-table--small .el-table__cell { + padding: var(--space-2) var(--space-3); +} + +.app-card .el-table--small th.el-table__cell { + font-size: 0.75rem; + font-weight: 600; + color: var(--muted-foreground); +} + +.app-card .el-table--small td.el-table__cell { + color: var(--el-text-color-regular); +} + +.app-card .el-table--small .el-table__row { + height: 40px; +} + +.app-card .el-table .el-button.is-link { + padding: 0; + height: auto; +} + +.app-card .el-table .cell { + line-height: 1.4; +} + +.table-scroll--sm { max-height: 320px; } +.table-scroll--md { max-height: 360px; } +.table-scroll--lg { max-height: 480px; } + +.card-toolbar .toolbar-input--grow { + flex: 1; + min-width: 200px; +} + +.card-pager .el-pagination { + font-size: 0.875rem; +} + +.kv-block p { + margin: 0 0 var(--space-2); +} + +.kv-block p:last-child { + margin-bottom: 0; +} + +.kv-block .label { + display: inline-block; + width: 5rem; + color: var(--muted-foreground); +} + +.kv-block code { + font-size: 0.8125rem; +} + +.cell-stack { + display: flex; + flex-direction: column; + gap: var(--space-2); + width: 100%; +} + +.meta-list { + margin: 0; + padding: 0; + list-style: none; + font-size: 0.8125rem; + color: var(--muted-foreground); +} + +.meta-list li { + display: flex; + justify-content: space-between; + gap: var(--space-3); + padding: var(--space-1) 0; +} + +.meta-list__mono { + font-family: ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, monospace; +} + +.meta-list__ok { color: var(--el-color-success); } +.meta-list__warn { color: var(--el-color-warning); } + +.chart-box { + height: 280px; +} + +.custom-scrollbar { + scrollbar-width: thin; + scrollbar-color: #e4e4e7 transparent; +} + +.custom-scrollbar::-webkit-scrollbar { + width: 6px; + height: 6px; +} + +.custom-scrollbar::-webkit-scrollbar-thumb { + background: #e4e4e7; + border-radius: 3px; +} + +.custom-scrollbar::-webkit-scrollbar-track { + background: transparent; +} + +@media (max-width: 1024px) { + .stat-grid { grid-template-columns: repeat(2, 1fr); } + .content-grid--2-1, + .content-grid--2, + .content-grid--3 { grid-template-columns: 1fr; } +} + +@media (max-width: 640px) { + .stat-grid { grid-template-columns: 1fr; } + .main-content { padding: var(--space-4); } + .page-header { flex-direction: column; } +} +``` + +--- + +*浅色后台管理系统 UI 设计规范* diff --git a/go.mod b/go.mod index e1f5b36..f4ad077 100644 --- a/go.mod +++ b/go.mod @@ -1,9 +1,10 @@ -module github.com/rose_cat707/go-vue-template +module github.com/rose_cat707/Atlas go 1.25.7 require ( github.com/gin-gonic/gin v1.12.0 + github.com/google/uuid v1.6.0 github.com/pressly/goose/v3 v3.27.1 golang.org/x/crypto v0.53.0 modernc.org/sqlite v1.52.0 @@ -22,7 +23,6 @@ require ( github.com/go-playground/validator/v10 v10.30.1 // indirect github.com/goccy/go-json v0.10.5 // indirect github.com/goccy/go-yaml v1.19.2 // indirect - github.com/google/uuid v1.6.0 // indirect github.com/json-iterator/go v1.1.12 // indirect github.com/klauspost/cpuid/v2 v2.3.0 // indirect github.com/leodido/go-urn v1.4.0 // indirect diff --git a/internal/api/admin_handlers.go b/internal/api/admin_handlers.go new file mode 100644 index 0000000..b30b569 --- /dev/null +++ b/internal/api/admin_handlers.go @@ -0,0 +1,199 @@ +package api + +import ( + "net/http" + "strconv" + + "github.com/gin-gonic/gin" + + "github.com/rose_cat707/Atlas/internal/store" +) + +func (s *Server) registerAdminRoutes(r gin.IRoutes) { + r.GET("/admin/providers", s.adminListProviders) + r.POST("/admin/providers", s.adminCreateProvider) + r.GET("/admin/providers/:id", s.adminGetProvider) + r.PUT("/admin/providers/:id", s.adminUpdateProvider) + r.DELETE("/admin/providers/:id", s.adminDeleteProvider) + r.GET("/admin/providers/:id/models", s.adminListProviderModels) + r.POST("/admin/providers/:id/models", s.adminAddProviderModel) + r.DELETE("/admin/providers/:id/models/:modelId", s.adminDeleteProviderModel) + r.POST("/admin/providers/:id/models/sync", s.adminSyncProviderModels) + r.GET("/admin/generations", s.adminListGenerations) + r.GET("/admin/generations/:id", s.adminGetGeneration) + r.DELETE("/admin/generations/:id", s.adminDeleteGeneration) + r.GET("/admin/dashboard/stats", s.adminDashboardStats) +} + +func (s *Server) adminListProviders(c *gin.Context) { + items, err := s.app.Store.ListProviders(c.Request.Context(), c.Query("service_type"), false) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + masked := make([]store.Provider, len(items)) + for i, p := range items { + masked[i] = store.MaskProvider(p) + } + c.JSON(http.StatusOK, gin.H{"items": masked}) +} + +func (s *Server) adminCreateProvider(c *gin.Context) { + var in store.Provider + if err := c.ShouldBindJSON(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := store.ValidateProvider(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if in.ExtraJSON == "" { + in.ExtraJSON = "{}" + } + id, err := s.app.Store.CreateProvider(c.Request.Context(), &in) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + in.ID = id + _ = s.app.Store.InsertAuditLog(c.Request.Context(), "create", "provider", in.Name) + c.JSON(http.StatusOK, store.MaskProvider(in)) +} + +func (s *Server) adminGetProvider(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + p, err := s.app.Store.GetProvider(c.Request.Context(), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + c.JSON(http.StatusOK, store.MaskProvider(*p)) +} + +func (s *Server) adminUpdateProvider(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + var in store.Provider + if err := c.ShouldBindJSON(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + existing, err := s.app.Store.GetProvider(c.Request.Context(), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + in.ID = id + if in.APIKey == "" || in.APIKey == "******" { + in.APIKey = existing.APIKey + } + if err := store.ValidateProvider(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := s.app.Store.UpdateProvider(c.Request.Context(), &in); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + _ = s.app.Store.InsertAuditLog(c.Request.Context(), "update", "provider", in.Name) + c.JSON(http.StatusOK, store.MaskProvider(in)) +} + +func (s *Server) adminDeleteProvider(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + if err := s.app.Store.DeleteProvider(c.Request.Context(), id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + _ = s.app.Store.InsertAuditLog(c.Request.Context(), "delete", "provider", strconv.FormatInt(id, 10)) + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (s *Server) adminListProviderModels(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + breakdown, err := s.app.Models.Breakdown(c.Request.Context(), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, breakdown) +} + +func (s *Server) adminAddProviderModel(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + var in struct { + ModelID string `json:"model_id"` + Name string `json:"name"` + } + if err := c.ShouldBindJSON(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + if err := s.app.Models.AddManual(c.Request.Context(), id, in.ModelID, in.Name); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + breakdown, _ := s.app.Models.Breakdown(c.Request.Context(), id) + c.JSON(http.StatusOK, breakdown) +} + +func (s *Server) adminDeleteProviderModel(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + modelID := c.Param("modelId") + if err := s.app.Models.RemoveManual(c.Request.Context(), id, modelID); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + breakdown, _ := s.app.Models.Breakdown(c.Request.Context(), id) + c.JSON(http.StatusOK, breakdown) +} + +func (s *Server) adminSyncProviderModels(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + upstream, err := s.app.Models.RefreshUpstream(c.Request.Context(), id) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + breakdown, _ := s.app.Models.Breakdown(c.Request.Context(), id) + c.JSON(http.StatusOK, gin.H{"upstream": upstream, "breakdown": breakdown}) +} + +func (s *Server) adminListGenerations(c *gin.Context) { + limit, offset := parseLimitOffset(c, 50, 500) + items, total, err := s.app.Store.ListAllGenerations(c.Request.Context(), c.Query("service_type"), limit, offset) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items, "total": total}) +} + +func (s *Server) adminGetGeneration(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + item, err := s.app.Store.AdminGetGeneration(c.Request.Context(), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + c.JSON(http.StatusOK, item) +} + +func (s *Server) adminDeleteGeneration(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + if err := s.app.Store.AdminDeleteGeneration(c.Request.Context(), id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + _ = s.app.Store.InsertAuditLog(c.Request.Context(), "delete", "generation", strconv.FormatInt(id, 10)) + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (s *Server) adminDashboardStats(c *gin.Context) { + stats, err := s.app.Store.DashboardStats(c.Request.Context()) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, stats) +} diff --git a/internal/api/gateway.go b/internal/api/gateway.go new file mode 100644 index 0000000..5e6be56 --- /dev/null +++ b/internal/api/gateway.go @@ -0,0 +1,25 @@ +package api + +import ( + "github.com/gin-gonic/gin" +) + +func (s *Server) registerGatewayRoutes(r *gin.Engine) { + v1 := r.Group("/v1") + v1.Use(s.sessionMiddleware()) + { + v1.GET("/models", s.gatewayModels) + + v1.POST("/chat/completions", s.gatewayChatCompletions) + v1.POST("/responses", s.gatewayResponses) + v1.POST("/completions", s.gatewayProxyOpenAI("/v1/completions")) + v1.POST("/embeddings", s.gatewayProxyOpenAI("/v1/embeddings")) + v1.POST("/rerank", s.gatewayProxyOpenAI("/v1/rerank")) + + v1.POST("/predictions", s.gatewayCreatePrediction) + v1.GET("/predictions", s.gatewayListPredictions) + v1.GET("/predictions/:id", s.gatewayGetPrediction) + v1.POST("/predictions/:id/cancel", s.gatewayCancelPrediction) + v1.GET("/predictions/:id/stream", s.gatewayStreamPrediction) + } +} diff --git a/internal/api/gateway_openai.go b/internal/api/gateway_openai.go new file mode 100644 index 0000000..bcc3a8f --- /dev/null +++ b/internal/api/gateway_openai.go @@ -0,0 +1,221 @@ +package api + +import ( + "encoding/json" + "fmt" + "io" + "net/http" + "strconv" + "strings" + + "github.com/gin-gonic/gin" + + "github.com/rose_cat707/Atlas/internal/domain" + "github.com/rose_cat707/Atlas/internal/provider" +) + +func (s *Server) gatewayProviderID(c *gin.Context) int64 { + if v := c.GetHeader("X-Provider-Id"); v != "" { + id, _ := strconv.ParseInt(v, 10, 64) + return id + } + return 0 +} + +func (s *Server) gatewayModels(c *gin.Context) { + ctx := c.Request.Context() + serviceType := c.DefaultQuery("service_type", string(domain.ServiceChat)) + protocol := c.DefaultQuery("protocol", domain.DefaultProtocol(domain.ServiceType(serviceType))) + cfg, err := s.app.Router.ResolveByProtocol(ctx, serviceType, protocol, s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + models, err := s.app.Models.ListEffective(ctx, cfg) + if err != nil || len(models) == 0 { + upstream, fetchErr := s.app.Models.FetchUpstream(ctx, cfg) + if fetchErr == nil && len(upstream) > 0 { + data := map[string]any{ + "object": "list", + "data": upstream, + } + out, _ := json.Marshal(data) + c.Data(http.StatusOK, "application/json", out) + return + } + path := provider.ModelsPath(cfg) + res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "GET", path, nil, nil, provider.CallMeta{ + SessionID: sessionID(c), + Method: "GET", + Path: path, + }) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.Data(res.StatusCode, "application/json", res.Body) + return + } + data := map[string]any{ + "object": "list", + "data": models, + } + out, _ := json.Marshal(data) + c.Data(http.StatusOK, "application/json", out) +} + +func (s *Server) gatewayChatCompletions(c *gin.Context) { + body, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + var req map[string]any + _ = json.Unmarshal(body, &req) + model, _ := req["model"].(string) + stream, _ := req["stream"].(bool) + + ctx := c.Request.Context() + cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceChat), "openai", s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + meta := provider.CallMeta{ + SessionID: sessionID(c), + Method: "POST", + Path: "/v1/chat/completions", + Model: model, + } + if stream { + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + err = s.app.Router.ProxyStream(ctx, cfg, "POST", "/v1/chat/completions", body, nil, meta, c.Writer, func() { c.Writer.Flush() }) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + } + return + } + res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "POST", "/v1/chat/completions", body, nil, meta) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.Data(res.StatusCode, "application/json", res.Body) +} + +func (s *Server) gatewayResponses(c *gin.Context) { + body, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + chatBody, model, err := responsesToChat(body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + ctx := c.Request.Context() + cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceChat), "openai", s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + meta := provider.CallMeta{ + SessionID: sessionID(c), + Method: "POST", + Path: "/v1/responses", + Model: model, + } + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + if err := s.app.Router.ProxyStream(ctx, cfg, "POST", "/v1/chat/completions", chatBody, nil, meta, c.Writer, func() { c.Writer.Flush() }); err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + } +} + +func (s *Server) gatewayProxyOpenAI(path string) gin.HandlerFunc { + return func(c *gin.Context) { + body, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + var req map[string]any + _ = json.Unmarshal(body, &req) + model, _ := req["model"].(string) + ctx := c.Request.Context() + cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceChat), "openai", s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + meta := provider.CallMeta{ + SessionID: sessionID(c), + Method: c.Request.Method, + Path: path, + Model: model, + } + res, err := s.app.Router.ProxyWithMeta(ctx, cfg, c.Request.Method, path, body, nil, meta) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.Data(res.StatusCode, "application/json", res.Body) + } +} + +func responsesToChat(body []byte) ([]byte, string, error) { + var req map[string]any + if err := json.Unmarshal(body, &req); err != nil { + return nil, "", err + } + model, _ := req["model"].(string) + messages := []map[string]string{} + if input, ok := req["input"].(string); ok && input != "" { + messages = append(messages, map[string]string{"role": "user", "content": input}) + } + if items, ok := req["input"].([]any); ok { + for _, item := range items { + m, ok := item.(map[string]any) + if !ok { + continue + } + role, _ := m["role"].(string) + if role == "" { + role = "user" + } + content, _ := m["content"].(string) + if content == "" { + if text, ok := m["text"].(string); ok { + content = text + } + } + if content != "" { + messages = append(messages, map[string]string{"role": role, "content": content}) + } + } + } + if len(messages) == 0 { + return nil, "", fmt.Errorf("responses input required") + } + out, err := json.Marshal(map[string]any{ + "model": model, + "messages": messages, + "stream": true, + }) + return out, model, err +} + +func inferImageServiceType(input map[string]any) domain.ServiceType { + if _, ok := input["image"]; ok { + return domain.ServiceImageToImage + } + return domain.ServiceTextToImage +} + +func gatewayPreferWait(c *gin.Context) bool { + return strings.Contains(strings.ToLower(c.GetHeader("Prefer")), "wait") +} diff --git a/internal/api/gateway_replicate.go b/internal/api/gateway_replicate.go new file mode 100644 index 0000000..f6d5c7c --- /dev/null +++ b/internal/api/gateway_replicate.go @@ -0,0 +1,153 @@ +package api + +import ( + "encoding/json" + "io" + "net/http" + + "github.com/gin-gonic/gin" + + "github.com/rose_cat707/Atlas/internal/domain" + "github.com/rose_cat707/Atlas/internal/provider" +) + +func (s *Server) gatewayCreatePrediction(c *gin.Context) { + body, err := io.ReadAll(c.Request.Body) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + var req map[string]any + if err := json.Unmarshal(body, &req); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + model, _ := req["model"].(string) + input, _ := req["input"].(map[string]any) + if input == nil { + input = map[string]any{} + } + params, _ := json.Marshal(input) + serviceType := inferImageServiceType(input) + mode := domain.ModeAsync + if gatewayPreferWait(c) { + mode = domain.ModeSync + } + ctx := c.Request.Context() + result, err := s.app.Invoke.Run(ctx, domain.InvokeRequest{ + SessionID: sessionID(c), + ServiceType: serviceType, + ProviderID: s.gatewayProviderID(c), + Model: model, + Params: params, + Mode: mode, + }) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + if mode == domain.ModeSync { + var output any + _ = json.Unmarshal(result.Output, &output) + c.JSON(http.StatusOK, gin.H{ + "id": "sync", + "status": "succeeded", + "output": output, + }) + return + } + task, err := s.app.Store.GetTask(ctx, sessionID(c), *result.TaskID) + if err != nil { + c.JSON(http.StatusOK, gin.H{"id": result.TaskID, "status": "starting"}) + return + } + c.JSON(http.StatusOK, gin.H{ + "id": task.ExternalID, + "status": task.Status, + }) +} + +func (s *Server) gatewayListPredictions(c *gin.Context) { + ctx := c.Request.Context() + cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + meta := provider.CallMeta{ + SessionID: sessionID(c), + Method: "GET", + Path: "/v1/predictions", + } + res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "GET", "/v1/predictions", nil, nil, meta) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.Data(res.StatusCode, "application/json", res.Body) +} + +func (s *Server) gatewayGetPrediction(c *gin.Context) { + ctx := c.Request.Context() + cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + path := "/v1/predictions/" + c.Param("id") + meta := provider.CallMeta{ + SessionID: sessionID(c), + Method: "GET", + Path: "/v1/predictions/:id", + } + res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "GET", path, nil, nil, meta) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.Data(res.StatusCode, "application/json", res.Body) +} + +func (s *Server) gatewayCancelPrediction(c *gin.Context) { + ctx := c.Request.Context() + id := c.Param("id") + cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + path := "/v1/predictions/" + id + "/cancel" + meta := provider.CallMeta{ + SessionID: sessionID(c), + Method: "POST", + Path: "/v1/predictions/:id/cancel", + } + res, err := s.app.Router.ProxyWithMeta(ctx, cfg, "POST", path, nil, nil, meta) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.Data(res.StatusCode, "application/json", res.Body) +} + +func (s *Server) gatewayStreamPrediction(c *gin.Context) { + ctx := c.Request.Context() + id := c.Param("id") + cfg, err := s.app.Router.ResolveByProtocol(ctx, string(domain.ServiceTextToImage), "replicate", s.gatewayProviderID(c)) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + path := "/v1/predictions/" + id + "/stream" + meta := provider.CallMeta{ + SessionID: sessionID(c), + Method: "GET", + Path: "/v1/predictions/:id/stream", + } + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + if err := s.app.Router.ProxyStream(ctx, cfg, "GET", path, nil, nil, meta, c.Writer, func() { c.Writer.Flush() }); err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + } +} diff --git a/internal/api/handlers.go b/internal/api/handlers.go index ded0ef4..d284342 100644 --- a/internal/api/handlers.go +++ b/internal/api/handlers.go @@ -9,16 +9,20 @@ import ( "github.com/gin-gonic/gin" "golang.org/x/crypto/bcrypt" - "github.com/rose_cat707/go-vue-template/internal/service" - "github.com/rose_cat707/go-vue-template/internal/settings" + "github.com/rose_cat707/Atlas/internal/service" + "github.com/rose_cat707/Atlas/internal/settings" ) type Server struct { - app *service.App + app *service.App + loginGuard *loginGuard } func NewServer(app *service.App) *Server { - return &Server{app: app} + return &Server{ + app: app, + loginGuard: newLoginGuard(), + } } func (s *Server) Router() *gin.Engine { @@ -41,8 +45,14 @@ func (s *Server) Router() *gin.Engine { protected.GET("/settings", s.getSettings) protected.PUT("/settings", s.putSettings) protected.GET("/audit-logs", s.listAuditLogs) + s.registerAdminRoutes(protected) } + user := api.Group("") + s.registerUserRoutes(user) + + s.registerGatewayRoutes(r) + r.NoRoute(s.serveSPA) return r } @@ -96,6 +106,15 @@ func (s *Server) authStatus(c *gin.Context) { } func (s *Server) login(c *gin.Context) { + ip := c.ClientIP() + if locked, retryAfter := s.loginGuard.check(ip); locked { + c.JSON(http.StatusTooManyRequests, gin.H{ + "error": "too many login attempts", + "retry_after": retryAfter, + }) + return + } + var req loginReq if err := c.ShouldBindJSON(&req); err != nil { c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) @@ -109,9 +128,18 @@ func (s *Server) login(c *gin.Context) { } hash := cfg[settings.KeyAdminPassword] if !checkPassword(req.Password, hash) { + time.Sleep(loginMinDelay) + if locked, retryAfter := s.loginGuard.recordFailure(ip); locked { + c.JSON(http.StatusTooManyRequests, gin.H{ + "error": "too many login attempts", + "retry_after": retryAfter, + }) + return + } c.JSON(http.StatusUnauthorized, gin.H{"error": "invalid password"}) return } + s.loginGuard.reset(ip) token := issueToken(hash, 24*time.Hour) c.JSON(http.StatusOK, gin.H{"token": token}) } diff --git a/internal/api/login_guard.go b/internal/api/login_guard.go new file mode 100644 index 0000000..0b0b4d5 --- /dev/null +++ b/internal/api/login_guard.go @@ -0,0 +1,68 @@ +package api + +import ( + "sync" + "time" +) + +const ( + loginMaxFailures = 5 + loginLockDuration = 15 * time.Minute + loginMinDelay = time.Second +) + +type loginGuard struct { + mu sync.Mutex + byIP map[string]*loginGuardEntry +} + +type loginGuardEntry struct { + failures int + lockedUntil time.Time +} + +func newLoginGuard() *loginGuard { + return &loginGuard{byIP: make(map[string]*loginGuardEntry)} +} + +func (g *loginGuard) check(ip string) (locked bool, retryAfter int) { + g.mu.Lock() + defer g.mu.Unlock() + + e := g.byIP[ip] + if e == nil { + return false, 0 + } + now := time.Now() + if now.Before(e.lockedUntil) { + return true, int(e.lockedUntil.Sub(now).Seconds()) + 1 + } + if e.failures >= loginMaxFailures { + e.failures = 0 + e.lockedUntil = time.Time{} + } + return false, 0 +} + +func (g *loginGuard) recordFailure(ip string) (locked bool, retryAfter int) { + g.mu.Lock() + defer g.mu.Unlock() + + e := g.byIP[ip] + if e == nil { + e = &loginGuardEntry{} + g.byIP[ip] = e + } + e.failures++ + if e.failures >= loginMaxFailures { + e.lockedUntil = time.Now().Add(loginLockDuration) + return true, int(loginLockDuration.Seconds()) + } + return false, 0 +} + +func (g *loginGuard) reset(ip string) { + g.mu.Lock() + delete(g.byIP, ip) + g.mu.Unlock() +} diff --git a/internal/api/login_guard_test.go b/internal/api/login_guard_test.go new file mode 100644 index 0000000..f3b5604 --- /dev/null +++ b/internal/api/login_guard_test.go @@ -0,0 +1,50 @@ +package api + +import ( + "testing" + "time" +) + +func TestLoginGuardLockout(t *testing.T) { + g := newLoginGuard() + ip := "203.0.113.1" + + for i := 0; i < loginMaxFailures-1; i++ { + locked, _ := g.recordFailure(ip) + if locked { + t.Fatalf("unexpected lock on attempt %d", i+1) + } + } + if locked, _ := g.check(ip); locked { + t.Fatal("should not be locked before threshold") + } + + locked, retryAfter := g.recordFailure(ip) + if !locked || retryAfter <= 0 { + t.Fatalf("expected lockout, locked=%v retryAfter=%d", locked, retryAfter) + } + if locked, _ := g.check(ip); !locked { + t.Fatal("expected locked after threshold") + } + + g.reset(ip) + if locked, _ := g.check(ip); locked { + t.Fatal("expected unlocked after reset") + } +} + +func TestLoginGuardLockExpires(t *testing.T) { + g := newLoginGuard() + ip := "203.0.113.2" + + g.mu.Lock() + g.byIP[ip] = &loginGuardEntry{ + failures: loginMaxFailures, + lockedUntil: time.Now().Add(-time.Second), + } + g.mu.Unlock() + + if locked, _ := g.check(ip); locked { + t.Fatal("expected expired lock to clear") + } +} diff --git a/internal/api/session.go b/internal/api/session.go new file mode 100644 index 0000000..83a4898 --- /dev/null +++ b/internal/api/session.go @@ -0,0 +1,40 @@ +package api + +import ( + "net/http" + "time" + + "github.com/gin-gonic/gin" + "github.com/google/uuid" +) + +const sessionCookie = "Atlas-Session" + +func (s *Server) sessionMiddleware() gin.HandlerFunc { + return func(c *gin.Context) { + sid, err := c.Cookie(sessionCookie) + if err != nil || sid == "" { + sid = uuid.New().String() + http.SetCookie(c.Writer, &http.Cookie{ + Name: sessionCookie, + Value: sid, + Path: "/", + MaxAge: int((30 * 24 * time.Hour).Seconds()), + HttpOnly: true, + SameSite: http.SameSiteLaxMode, + }) + } + _ = s.app.Store.UpsertSession(c.Request.Context(), sid) + c.Set("session_id", sid) + c.Next() + } +} + +func sessionID(c *gin.Context) string { + if v, ok := c.Get("session_id"); ok { + if s, ok := v.(string); ok { + return s + } + } + return "" +} diff --git a/internal/api/task_sse.go b/internal/api/task_sse.go new file mode 100644 index 0000000..4480226 --- /dev/null +++ b/internal/api/task_sse.go @@ -0,0 +1,84 @@ +package api + +import ( + "context" + "encoding/json" + "fmt" + "net/http" + "strconv" + "time" + + "github.com/gin-gonic/gin" + + "github.com/rose_cat707/Atlas/internal/service/taskevents" + "github.com/rose_cat707/Atlas/internal/store" +) + +func (s *Server) streamTaskEvents(c *gin.Context) { + id, err := strconv.ParseInt(c.Param("id"), 10, 64) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "invalid task id"}) + return + } + ctx := c.Request.Context() + sid := sessionID(c) + task, err := s.resolveTask(ctx, sid, id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Status(http.StatusOK) + + writeTaskSSE(c, task) + if taskevents.Terminal(task.Status) { + writeSSEDone(c) + return + } + + ch := s.app.TaskEvents.Subscribe(id) + defer s.app.TaskEvents.Unsubscribe(id, ch) + + heartbeat := time.NewTicker(15 * time.Second) + defer heartbeat.Stop() + + for { + select { + case <-ctx.Done(): + return + case t := <-ch: + writeTaskSSE(c, &t) + if taskevents.Terminal(t.Status) { + writeSSEDone(c) + return + } + case <-heartbeat.C: + fmt.Fprintf(c.Writer, ": heartbeat\n\n") + c.Writer.Flush() + } + } +} + +func (s *Server) resolveTask(ctx context.Context, sessionID string, id int64) (*store.Task, error) { + key := "task:" + strconv.FormatInt(id, 10) + if v, ok := s.app.Store.Cache.Get(key); ok { + if t, ok := v.(*store.Task); ok { + return t, nil + } + } + return s.app.Store.GetTask(ctx, sessionID, id) +} + +func writeTaskSSE(c *gin.Context, task *store.Task) { + data, _ := json.Marshal(task) + fmt.Fprintf(c.Writer, "event: task\ndata: %s\n\n", data) + c.Writer.Flush() +} + +func writeSSEDone(c *gin.Context) { + fmt.Fprintf(c.Writer, "event: done\ndata: [DONE]\n\n") + c.Writer.Flush() +} diff --git a/internal/api/user_handlers.go b/internal/api/user_handlers.go new file mode 100644 index 0000000..298678c --- /dev/null +++ b/internal/api/user_handlers.go @@ -0,0 +1,427 @@ +package api + +import ( + "encoding/json" + "net/http" + "os" + "strconv" + "time" + + "github.com/gin-gonic/gin" + + "github.com/rose_cat707/Atlas/internal/domain" + "github.com/rose_cat707/Atlas/internal/provider" + "github.com/rose_cat707/Atlas/internal/service/asset" + "github.com/rose_cat707/Atlas/internal/service/chat" + "github.com/rose_cat707/Atlas/internal/store" +) + +func (s *Server) registerUserRoutes(r gin.IRoutes) { + r.Use(s.sessionMiddleware()) + r.POST("/user/session", s.userSession) + r.GET("/user/meta", s.userMeta) + r.GET("/user/conversations", s.listConversations) + r.POST("/user/conversations", s.createConversation) + r.POST("/user/conversations/messages", s.postNewConversationMessage) + r.GET("/user/conversations/:id", s.getConversation) + r.PUT("/user/conversations/:id", s.updateConversation) + r.DELETE("/user/conversations/:id", s.deleteConversation) + r.GET("/user/conversations/:id/messages", s.listMessages) + r.POST("/user/conversations/:id/messages", s.postMessage) + r.POST("/user/generate", s.userGenerate) + r.POST("/user/upload", s.userUpload) + r.GET("/user/files/:id", s.serveFile) + r.GET("/user/generations", s.listGenerations) + r.GET("/user/generations/:id", s.getGeneration) + r.DELETE("/user/generations/:id", s.deleteGeneration) + r.GET("/user/presets", s.listPresets) + r.POST("/user/presets", s.createPreset) + r.PUT("/user/presets/:id", s.updatePreset) + r.DELETE("/user/presets/:id", s.deletePreset) + r.GET("/user/tasks/active", s.listActiveTasks) + r.GET("/user/tasks/:id/events", s.streamTaskEvents) + r.GET("/user/tasks/:id", s.getTask) + r.POST("/user/tasks/:id/cancel", s.cancelTask) +} + +func (s *Server) userSession(c *gin.Context) { + c.JSON(http.StatusOK, gin.H{"session_id": sessionID(c)}) +} + +func (s *Server) userMeta(c *gin.Context) { + ctx := c.Request.Context() + specs := s.app.Invoke.Specs() + imageModels, _ := s.app.Models.ListImageModelsForMeta(ctx) + out := gin.H{} + for st, spec := range specs { + providers, err := s.app.Store.ListProviders(ctx, string(st), true) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + var provList []gin.H + var allModels any + switch st { + case domain.ServiceTextToImage: + filtered := provider.FilterModelsByKind(imageModels, provider.ModelKindTextToImage) + allModels = filtered + case domain.ServiceImageToImage: + filtered := provider.FilterModelsByKind(imageModels, provider.ModelKindImageToImage) + allModels = filtered + default: + var entries []provider.ModelEntry + for _, p := range providers { + cfg, err := s.app.Router.Resolve(c.Request.Context(), string(st), p.ID) + if err != nil { + continue + } + models, err := s.app.Models.ListEffective(ctx, cfg) + if err != nil { + continue + } + entries = append(entries, models...) + } + allModels = entries + } + for _, p := range providers { + provList = append(provList, gin.H{"id": p.ID, "name": p.Name, "is_default": p.IsDefault}) + } + caps := gin.H{} + if spec.Capabilities.Stream != nil { + caps["stream"] = true + } + if spec.Capabilities.Sync != nil { + caps["sync"] = true + } + if spec.Capabilities.Async != nil { + caps["async"] = true + } + defaultModel := "" + switch models := allModels.(type) { + case []provider.ModelDetail: + defaultModel = provider.DefaultModelDetail(models) + case []provider.ModelEntry: + defaultModel = provider.DefaultModel(models) + } + out[string(st)] = gin.H{ + "capabilities": caps, + "providers": provList, + "models": allModels, + "default_model": defaultModel, + } + } + c.JSON(http.StatusOK, gin.H{"service_types": out}) +} + +func (s *Server) listConversations(c *gin.Context) { + items, err := s.app.Store.ListConversations(c.Request.Context(), sessionID(c)) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items}) +} + +func (s *Server) createConversation(c *gin.Context) { + c.JSON(http.StatusBadRequest, gin.H{"error": "conversation is created when sending the first message"}) +} + +func (s *Server) postNewConversationMessage(c *gin.Context) { + s.streamConversationMessage(c, 0) +} + +func (s *Server) getConversation(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + conv, err := s.app.Store.GetConversation(c.Request.Context(), sessionID(c), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + c.JSON(http.StatusOK, conv) +} + +func (s *Server) updateConversation(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + var in store.Conversation + if err := c.ShouldBindJSON(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + in.ID = id + in.SessionID = sessionID(c) + if err := s.app.Store.UpdateConversation(c.Request.Context(), &in); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, in) +} + +func (s *Server) deleteConversation(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + if err := s.app.Store.DeleteConversation(c.Request.Context(), sessionID(c), id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (s *Server) listMessages(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + items, err := s.app.Store.ListMessages(c.Request.Context(), id, 80) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items}) +} + +func (s *Server) postMessage(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + s.streamConversationMessage(c, id) +} + +func (s *Server) streamConversationMessage(c *gin.Context, convID int64) { + var in struct { + Content string `json:"content"` + Model string `json:"model"` + ProviderID int64 `json:"provider_id"` + } + if err := c.ShouldBindJSON(&in); err != nil || in.Content == "" { + c.JSON(http.StatusBadRequest, gin.H{"error": "content required"}) + return + } + + ctx := c.Request.Context() + sid := sessionID(c) + if convID == 0 { + conv := &store.Conversation{ + SessionID: sid, + Title: "新对话", + Model: in.Model, + ParamsJSON: "{}", + } + id, err := s.app.Store.CreateConversation(ctx, conv) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + convID = id + } else if _, err := s.app.Store.GetConversation(ctx, sid, convID); err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + + c.Header("X-Conversation-Id", strconv.FormatInt(convID, 10)) + c.Header("Content-Type", "text/event-stream") + c.Header("Cache-Control", "no-cache") + c.Header("Connection", "keep-alive") + c.Status(http.StatusOK) + err := s.app.Chat.StreamMessage(ctx, chat.StreamRequest{ + SessionID: sid, + ConversationID: convID, + Content: in.Content, + Model: in.Model, + ProviderID: in.ProviderID, + }, c.Writer, func() { c.Writer.Flush() }) + if err != nil { + _, _ = c.Writer.Write([]byte("data: " + `{"error":"` + err.Error() + `"}` + "\n\n")) + c.Writer.Flush() + } +} + +func (s *Server) userGenerate(c *gin.Context) { + var in struct { + ServiceType string `json:"service_type"` + Params json.RawMessage `json:"params"` + Mode string `json:"mode"` + Model string `json:"model"` + ProviderID int64 `json:"provider_id"` + } + if err := c.ShouldBindJSON(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + mode := domain.InvokeMode(in.Mode) + if mode == "" { + mode = domain.ModeSync + } + result, err := s.app.Invoke.Run(c.Request.Context(), domain.InvokeRequest{ + SessionID: sessionID(c), + ServiceType: domain.ServiceType(in.ServiceType), + ProviderID: in.ProviderID, + Model: in.Model, + Params: in.Params, + Mode: mode, + }) + if err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, result) +} + +func (s *Server) userUpload(c *gin.Context) { + file, err := c.FormFile("file") + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": "file required"}) + return + } + f, err := file.Open() + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + defer f.Close() + path, size, mimeType, err := asset.SaveUpload(s.app.DataDir, sessionID(c), file.Filename, f, 10<<20) + if err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + exp := time.Now().Add(7 * 24 * time.Hour).UTC().Format(time.RFC3339) + rec := &store.FileRecord{ + SessionID: sessionID(c), + Kind: "upload", + Mime: mimeType, + StoragePath: path, + SizeBytes: size, + ExpiresAt: &exp, + } + id, err := s.app.Store.InsertFile(c.Request.Context(), rec) + if err != nil { + os.Remove(path) + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + signed := s.app.Assets.SignedURL("/api/v1", id, 7*24*time.Hour) + c.JSON(http.StatusOK, gin.H{"file_id": id, "signed_url": signed, "mime": mimeType}) +} + +func (s *Server) serveFile(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + expStr := c.Query("exp") + sig := c.Query("sig") + exp, sig, err := asset.ParseSignedQuery(expStr, sig) + if err != nil || !s.app.Assets.Verify(id, exp, sig) { + c.JSON(http.StatusForbidden, gin.H{"error": "invalid signature"}) + return + } + rec, err := s.app.Store.GetFile(c.Request.Context(), sessionID(c), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + c.Header("Content-Type", rec.Mime) + c.File(rec.StoragePath) +} + +func (s *Server) listGenerations(c *gin.Context) { + limit, offset := parseLimitOffset(c, 20, 100) + st := c.Query("service_type") + items, total, err := s.app.Store.ListGenerations(c.Request.Context(), sessionID(c), st, limit, offset) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items, "total": total}) +} + +func (s *Server) getGeneration(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + g, err := s.app.Store.GetGeneration(c.Request.Context(), sessionID(c), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + c.JSON(http.StatusOK, g) +} + +func (s *Server) deleteGeneration(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + if err := s.app.Store.DeleteGeneration(c.Request.Context(), sessionID(c), id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (s *Server) listPresets(c *gin.Context) { + items, err := s.app.Store.ListPresets(c.Request.Context(), sessionID(c), c.Query("service_type")) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items}) +} + +func (s *Server) createPreset(c *gin.Context) { + var in store.Preset + if err := c.ShouldBindJSON(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + in.SessionID = sessionID(c) + id, err := s.app.Store.CreatePreset(c.Request.Context(), &in) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + in.ID = id + c.JSON(http.StatusOK, in) +} + +func (s *Server) updatePreset(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + var in store.Preset + if err := c.ShouldBindJSON(&in); err != nil { + c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()}) + return + } + in.ID = id + in.SessionID = sessionID(c) + if err := s.app.Store.UpdatePreset(c.Request.Context(), &in); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, in) +} + +func (s *Server) deletePreset(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + if err := s.app.Store.DeletePreset(c.Request.Context(), sessionID(c), id); err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} + +func (s *Server) listActiveTasks(c *gin.Context) { + items, err := s.app.Store.ListActiveTasks(c.Request.Context(), sessionID(c), c.Query("service_type")) + if err != nil { + c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"items": items}) +} + +func (s *Server) getTask(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + if v, ok := s.app.Store.Cache.Get("task:" + strconv.FormatInt(id, 10)); ok { + c.JSON(http.StatusOK, v) + return + } + t, err := s.app.Store.GetTask(c.Request.Context(), sessionID(c), id) + if err != nil { + c.JSON(http.StatusNotFound, gin.H{"error": "not found"}) + return + } + c.JSON(http.StatusOK, t) +} + +func (s *Server) cancelTask(c *gin.Context) { + id, _ := strconv.ParseInt(c.Param("id"), 10, 64) + if err := s.app.Invoke.CancelTask(c.Request.Context(), sessionID(c), id); err != nil { + c.JSON(http.StatusBadGateway, gin.H{"error": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"ok": true}) +} diff --git a/internal/domain/invoke.go b/internal/domain/invoke.go new file mode 100644 index 0000000..3c28c7c --- /dev/null +++ b/internal/domain/invoke.go @@ -0,0 +1,52 @@ +package domain + +import "encoding/json" + +type ServiceType string + +const ( + ServiceChat ServiceType = "chat" + ServiceTextToImage ServiceType = "text_to_image" + ServiceImageToImage ServiceType = "image_to_image" +) + +type InvokeMode string + +const ( + ModeStream InvokeMode = "stream" + ModeSync InvokeMode = "sync" + ModeAsync InvokeMode = "async" +) + +type InvokeContext struct { + ConversationID int64 `json:"conversation_id,omitempty"` + Messages []ChatMessage `json:"messages,omitempty"` +} + +type ChatMessage struct { + Role string `json:"role"` + Content string `json:"content"` +} + +type InvokeRequest struct { + SessionID string `json:"session_id"` + ServiceType ServiceType `json:"service_type"` + ProviderID int64 `json:"provider_id,omitempty"` + Model string `json:"model"` + Params json.RawMessage `json:"params"` + Mode InvokeMode `json:"mode"` + Context *InvokeContext `json:"context,omitempty"` +} + +type StreamChunk struct { + Delta string `json:"delta"` + Done bool `json:"done"` + Error string `json:"error,omitempty"` +} + +type InvokeResult struct { + GenerationID int64 `json:"generation_id,omitempty"` + TaskID *int64 `json:"task_id,omitempty"` + Output json.RawMessage `json:"output,omitempty"` + Stream <-chan StreamChunk `json:"-"` +} diff --git a/internal/domain/provider_compat.go b/internal/domain/provider_compat.go new file mode 100644 index 0000000..b664e6f --- /dev/null +++ b/internal/domain/provider_compat.go @@ -0,0 +1,51 @@ +package domain + +import "fmt" + +var protocolServiceTypes = map[string][]ServiceType{ + "openai": {ServiceChat}, + "replicate": {ServiceTextToImage, ServiceImageToImage}, +} + +var serviceTypeProtocol = map[ServiceType]string{ + ServiceChat: "openai", + ServiceTextToImage: "replicate", + ServiceImageToImage: "replicate", +} + +func DefaultProtocol(serviceType ServiceType) string { + if p, ok := serviceTypeProtocol[serviceType]; ok { + return p + } + return "openai" +} + +func DefaultServiceType(protocol string) ServiceType { + types := ServiceTypesForProtocol(protocol) + if len(types) == 0 { + return ServiceChat + } + return types[0] +} + +func ServiceTypesForProtocol(protocol string) []ServiceType { + return append([]ServiceType(nil), protocolServiceTypes[protocol]...) +} + +func ProtocolForServiceType(serviceType ServiceType) string { + return DefaultProtocol(serviceType) +} + +func ValidateProviderBinding(protocol, serviceType string) error { + types := ServiceTypesForProtocol(protocol) + if len(types) == 0 { + return fmt.Errorf("unsupported protocol: %s", protocol) + } + st := ServiceType(serviceType) + for _, allowed := range types { + if allowed == st { + return nil + } + } + return fmt.Errorf("protocol %s does not support service_type %s", protocol, serviceType) +} diff --git a/internal/domain/provider_compat_test.go b/internal/domain/provider_compat_test.go new file mode 100644 index 0000000..ab6c29d --- /dev/null +++ b/internal/domain/provider_compat_test.go @@ -0,0 +1,36 @@ +package domain + +import "testing" + +func TestValidateProviderBinding(t *testing.T) { + cases := []struct { + protocol string + serviceType string + ok bool + }{ + {"openai", "chat", true}, + {"replicate", "text_to_image", true}, + {"replicate", "image_to_image", true}, + {"openai", "text_to_image", false}, + {"replicate", "chat", false}, + {"unknown", "chat", false}, + } + for _, tc := range cases { + err := ValidateProviderBinding(tc.protocol, tc.serviceType) + if tc.ok && err != nil { + t.Fatalf("%s/%s expected ok, got %v", tc.protocol, tc.serviceType, err) + } + if !tc.ok && err == nil { + t.Fatalf("%s/%s expected error", tc.protocol, tc.serviceType) + } + } +} + +func TestDefaultBindings(t *testing.T) { + if got := DefaultProtocol(ServiceTextToImage); got != "replicate" { + t.Fatalf("expected replicate, got %s", got) + } + if got := DefaultServiceType("openai"); got != ServiceChat { + t.Fatalf("expected chat, got %s", got) + } +} diff --git a/internal/domain/servicetype.go b/internal/domain/servicetype.go new file mode 100644 index 0000000..bc20ab7 --- /dev/null +++ b/internal/domain/servicetype.go @@ -0,0 +1,30 @@ +package domain + +import "time" + +type StreamCapability struct{} + +type SyncCapability struct{} + +type AsyncCapability struct { + Driver string + PollInterval time.Duration +} + +type InvokeCapabilities struct { + Stream *StreamCapability + Sync *SyncCapability + Async *AsyncCapability +} + +type ModelInfo struct { + ID string `json:"id"` + Name string `json:"name"` +} + +type ServiceTypeSpec struct { + Type ServiceType + Capabilities InvokeCapabilities + DefaultParams func() map[string]any + ValidateParams func(params map[string]any) error +} diff --git a/internal/provider/model_schema.go b/internal/provider/model_schema.go new file mode 100644 index 0000000..e130530 --- /dev/null +++ b/internal/provider/model_schema.go @@ -0,0 +1,245 @@ +package provider + +import ( + "encoding/json" + "sort" + "strings" +) + +const ( + ModelKindTextToImage = "text_to_image" + ModelKindImageToImage = "image_to_image" +) + +type InputField struct { + Name string `json:"name"` + Title string `json:"title"` + Description string `json:"description"` + Type string `json:"type"` + Default any `json:"default,omitempty"` + Required bool `json:"required"` + Order int `json:"order"` + Widget string `json:"widget"` // text | textarea | number | integer | image | switch +} + +type ModelDetail struct { + ID string `json:"id"` + Name string `json:"name"` + Description string `json:"description,omitempty"` + Kind string `json:"kind"` + InputSchema []InputField `json:"input_schema"` +} + +type rawProperty struct { + Type string `json:"type"` + Title string `json:"title"` + Description string `json:"description"` + Default any `json:"default"` + Order int `json:"x-order"` +} + +func parseReplicateModelDetails(body []byte) []ModelDetail { + var resp struct { + Results []struct { + Owner string `json:"owner"` + Name string `json:"name"` + Description string `json:"description"` + LatestVersion struct { + OpenAPISchema struct { + Components struct { + Schemas struct { + Input struct { + Properties map[string]json.RawMessage `json:"properties"` + Required []string `json:"required"` + } `json:"Input"` + } `json:"schemas"` + } `json:"components"` + } `json:"openapi_schema"` + } `json:"latest_version"` + } `json:"results"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return nil + } + out := make([]ModelDetail, 0, len(resp.Results)) + for _, item := range resp.Results { + if item.Owner == "" || item.Name == "" { + continue + } + id := item.Owner + "/" + item.Name + name := item.Name + if item.Description != "" { + name = item.Description + } + schema := parseInputSchema(item.LatestVersion.OpenAPISchema.Components.Schemas.Input.Properties, item.LatestVersion.OpenAPISchema.Components.Schemas.Input.Required) + out = append(out, ModelDetail{ + ID: id, + Name: name, + Description: item.Description, + Kind: inferModelKind(schema), + InputSchema: schema, + }) + } + return out +} + +func parseInputSchema(properties map[string]json.RawMessage, required []string) []InputField { + if len(properties) == 0 { + return defaultPromptSchema() + } + reqSet := map[string]struct{}{} + for _, name := range required { + reqSet[name] = struct{}{} + } + fields := make([]InputField, 0, len(properties)) + for name, raw := range properties { + var prop rawProperty + if err := json.Unmarshal(raw, &prop); err != nil { + continue + } + title := prop.Title + if title == "" { + title = name + } + fieldType := prop.Type + if fieldType == "" { + fieldType = "string" + } + fields = append(fields, InputField{ + Name: name, + Title: title, + Description: prop.Description, + Type: fieldType, + Default: prop.Default, + Required: containsRequired(reqSet, name), + Order: prop.Order, + Widget: inferWidget(name, fieldType), + }) + } + sort.Slice(fields, func(i, j int) bool { + if fields[i].Order == fields[j].Order { + return fields[i].Name < fields[j].Name + } + return fields[i].Order < fields[j].Order + }) + return fields +} + +func inferModelKind(schema []InputField) string { + for _, f := range schema { + if f.Name == "image" && f.Required { + return ModelKindImageToImage + } + } + for _, f := range schema { + if f.Name == "image" && f.Widget == "image" { + return ModelKindImageToImage + } + } + return ModelKindTextToImage +} + +func inferWidget(name, fieldType string) string { + switch fieldType { + case "integer", "number": + return fieldType + case "boolean": + return "switch" + } + if name == "image" || strings.HasSuffix(name, "_image") || strings.Contains(name, "image_url") { + return "image" + } + if name == "prompt" || name == "negative_prompt" || strings.Contains(name, "prompt") { + return "textarea" + } + return "text" +} + +func defaultPromptSchema() []InputField { + return []InputField{{ + Name: "prompt", + Title: "正向提示词", + Type: "string", + Required: true, + Order: 0, + Widget: "textarea", + }} +} + +func manualModelDetail(modelID, name string) ModelDetail { + if name == "" { + name = modelID + } + schema := defaultPromptSchema() + kind := ModelKindTextToImage + if strings.Contains(strings.ToLower(modelID), "edit") || strings.Contains(strings.ToLower(name), "edit") { + schema = append(schema, InputField{ + Name: "image", + Title: "输入图 URL", + Type: "string", + Required: true, + Order: 2, + Widget: "image", + }) + kind = ModelKindImageToImage + } + return ModelDetail{ + ID: modelID, + Name: name, + Kind: kind, + InputSchema: schema, + } +} + +func containsRequired(set map[string]struct{}, name string) bool { + _, ok := set[name] + return ok +} + +func ValidateModelInput(schema []InputField, params map[string]any) error { + if len(schema) == 0 { + if p, _ := params["prompt"].(string); strings.TrimSpace(p) == "" { + return errRequired("prompt") + } + return nil + } + for _, field := range schema { + if !field.Required { + continue + } + val, ok := params[field.Name] + if !ok || isEmptyValue(val) { + return errRequired(field.Name) + } + } + return nil +} + +func isEmptyValue(v any) bool { + switch t := v.(type) { + case string: + return strings.TrimSpace(t) == "" + case nil: + return true + default: + return false + } +} + +type validationError string + +func (e validationError) Error() string { return string(e) } + +func errRequired(name string) error { + return validationError(name + " required") +} + +func DefaultParamValues(schema []InputField) map[string]any { + out := map[string]any{} + for _, field := range schema { + if field.Default != nil { + out[field.Name] = field.Default + } + } + return out +} diff --git a/internal/provider/model_schema_test.go b/internal/provider/model_schema_test.go new file mode 100644 index 0000000..36a0909 --- /dev/null +++ b/internal/provider/model_schema_test.go @@ -0,0 +1,43 @@ +package provider + +import ( + "os" + "testing" +) + +func TestParseReplicateModelDetails(t *testing.T) { + body, err := os.ReadFile("testdata/models_sample.json") + if err != nil { + t.Skip("sample file missing") + } + models := parseReplicateModelDetails(body) + if len(models) != 2 { + t.Fatalf("expected 2 models, got %d", len(models)) + } + if models[0].Kind != ModelKindImageToImage { + t.Fatalf("edit model kind: %s", models[0].Kind) + } + if models[1].Kind != ModelKindTextToImage { + t.Fatalf("text model kind: %s", models[1].Kind) + } + if len(models[0].InputSchema) < 3 { + t.Fatalf("expected schema fields") + } +} + +func TestInferModelKind(t *testing.T) { + schema := []InputField{{Name: "prompt", Required: true}, {Name: "image", Required: true, Widget: "image"}} + if inferModelKind(schema) != ModelKindImageToImage { + t.Fatal("expected image_to_image") + } +} + +func TestValidateModelInput(t *testing.T) { + schema := []InputField{{Name: "prompt", Required: true}} + if err := ValidateModelInput(schema, map[string]any{"prompt": "hi"}); err != nil { + t.Fatal(err) + } + if err := ValidateModelInput(schema, map[string]any{}); err == nil { + t.Fatal("expected error") + } +} diff --git a/internal/provider/models.go b/internal/provider/models.go new file mode 100644 index 0000000..1adb3c2 --- /dev/null +++ b/internal/provider/models.go @@ -0,0 +1,432 @@ +package provider + +import ( + "context" + "encoding/json" + "fmt" + "net/url" + "strings" + + "github.com/rose_cat707/Atlas/internal/store" +) + +type ModelEntry struct { + ID string `json:"id"` + Name string `json:"name"` +} + +type ModelsService struct { + store *store.Store + router *Router +} + +func NewModelsService(st *store.Store, router *Router) *ModelsService { + return &ModelsService{store: st, router: router} +} + +func DefaultModel(models []ModelEntry) string { + if len(models) > 0 { + return models[0].ID + } + return "" +} + +// ListEffective returns upstream models first, then manual models not already present. +func (m *ModelsService) ListEffective(ctx context.Context, cfg *Config) ([]ModelEntry, error) { + upstream, _ := m.fetchUpstream(ctx, cfg) + manual, err := m.listManual(ctx, cfg.ID) + if err != nil { + return nil, err + } + return mergeModels(upstream, manual), nil +} + +func (m *ModelsService) ListForServiceType(ctx context.Context, serviceType string) ([]ModelEntry, error) { + providers, err := m.store.ListProviders(ctx, serviceType, true) + if err != nil { + return nil, err + } + seen := map[string]struct{}{} + var out []ModelEntry + for _, p := range providers { + cfg, err := m.router.providerToConfig(&p) + if err != nil { + continue + } + models, err := m.ListEffective(ctx, cfg) + if err != nil { + continue + } + for _, model := range models { + if _, ok := seen[model.ID]; ok { + continue + } + seen[model.ID] = struct{}{} + out = append(out, model) + } + } + return out, nil +} + +type ModelsBreakdown struct { + Upstream []ModelEntry `json:"upstream"` + Manual []ModelEntry `json:"manual"` + Effective []ModelEntry `json:"effective"` +} + +func (m *ModelsService) Breakdown(ctx context.Context, providerID int64) (*ModelsBreakdown, error) { + p, err := m.store.GetProvider(ctx, providerID) + if err != nil { + return nil, err + } + cfg, err := m.router.providerToConfig(p) + if err != nil { + return nil, err + } + upstream, _ := m.fetchUpstream(ctx, cfg) + manual, err := m.listManual(ctx, providerID) + if err != nil { + return nil, err + } + return &ModelsBreakdown{ + Upstream: upstream, + Manual: manual, + Effective: mergeModels(upstream, manual), + }, nil +} + +func (m *ModelsService) AddManual(ctx context.Context, providerID int64, modelID, name string) error { + modelID = strings.TrimSpace(modelID) + name = strings.TrimSpace(name) + if modelID == "" || name == "" { + return fmt.Errorf("model_id and name required") + } + _, err := m.store.CreateProviderModel(ctx, &store.ProviderModel{ + ProviderID: providerID, + ModelID: modelID, + Name: name, + }) + if err != nil { + return fmt.Errorf("model already exists or invalid") + } + m.invalidateCache(providerID) + return nil +} + +func (m *ModelsService) RemoveManual(ctx context.Context, providerID int64, modelID string) error { + if err := m.store.DeleteProviderModel(ctx, providerID, modelID); err != nil { + return err + } + m.invalidateCache(providerID) + return nil +} + +func (m *ModelsService) RefreshUpstream(ctx context.Context, providerID int64) ([]ModelEntry, error) { + p, err := m.store.GetProvider(ctx, providerID) + if err != nil { + return nil, err + } + cfg, err := m.router.providerToConfig(p) + if err != nil { + return nil, err + } + m.invalidateCache(providerID) + return m.fetchUpstream(ctx, cfg) +} + +// FetchUpstream loads models from upstream using the provider protocol. +func (m *ModelsService) FetchUpstream(ctx context.Context, cfg *Config) ([]ModelEntry, error) { + return m.fetchUpstream(ctx, cfg) +} + +// ModelsPath returns the upstream models API path for a provider configuration. +func ModelsPath(cfg *Config) string { + q := url.Values{} + if cfg.Protocol != "" { + q.Set("protocol", cfg.Protocol) + } + if cfg.ServiceType != "" { + q.Set("service_type", cfg.ServiceType) + } + if len(q) == 0 { + return "/v1/models" + } + return "/v1/models?" + q.Encode() +} + +func (m *ModelsService) listManual(ctx context.Context, providerID int64) ([]ModelEntry, error) { + rows, err := m.store.ListProviderModels(ctx, providerID) + if err != nil { + return nil, err + } + out := make([]ModelEntry, 0, len(rows)) + for _, r := range rows { + out = append(out, ModelEntry{ID: r.ModelID, Name: r.Name}) + } + return out, nil +} + +// ListEffectiveDetails returns upstream + manual models with input schema. +func (m *ModelsService) ListEffectiveDetails(ctx context.Context, cfg *Config) ([]ModelDetail, error) { + upstream, _ := m.fetchUpstreamDetails(ctx, cfg) + manual, err := m.listManual(ctx, cfg.ID) + if err != nil { + return nil, err + } + return mergeModelDetails(upstream, manual), nil +} + +// FindModelDetail looks up a model by id for a provider configuration. +func (m *ModelsService) FindModelDetail(ctx context.Context, cfg *Config, modelID string) (*ModelDetail, error) { + models, err := m.ListEffectiveDetails(ctx, cfg) + if err != nil { + return nil, err + } + for _, md := range models { + if md.ID == modelID { + copy := md + return ©, nil + } + } + return nil, fmt.Errorf("model not found: %s", modelID) +} + +// ListImageModelsForMeta aggregates image models from all image-related providers. +func (m *ModelsService) ListImageModelsForMeta(ctx context.Context) ([]ModelDetail, error) { + seen := map[string]struct{}{} + var out []ModelDetail + for _, serviceType := range []string{ModelKindTextToImage, ModelKindImageToImage} { + providers, err := m.store.ListProviders(ctx, serviceType, true) + if err != nil { + return nil, err + } + for _, p := range providers { + cfg, err := m.router.providerToConfig(&p) + if err != nil { + continue + } + models, err := m.ListEffectiveDetails(ctx, cfg) + if err != nil { + continue + } + for _, model := range models { + if _, ok := seen[model.ID]; ok { + continue + } + seen[model.ID] = struct{}{} + out = append(out, model) + } + } + } + return out, nil +} + +func (m *ModelsService) fetchUpstream(ctx context.Context, cfg *Config) ([]ModelEntry, error) { + details, err := m.fetchUpstreamDetails(ctx, cfg) + if err != nil || len(details) == 0 { + return nil, err + } + out := make([]ModelEntry, len(details)) + for i, d := range details { + out[i] = ModelEntry{ID: d.ID, Name: d.Name} + } + return out, nil +} + +func (m *ModelsService) fetchUpstreamDetails(ctx context.Context, cfg *Config) ([]ModelDetail, error) { + key := detailsCacheKey(cfg.ID) + if v, ok := m.store.Cache.Get(key); ok { + if models, ok := v.([]ModelDetail); ok { + return models, nil + } + } + path := ModelsPath(cfg) + meta := CallMeta{Method: "GET", Path: path} + res, err := m.router.ProxyWithMeta(ctx, cfg, "GET", path, nil, nil, meta) + if err != nil || res.StatusCode >= 400 { + return nil, err + } + models := parseDetailsForProtocol(cfg.Protocol, res.Body) + if len(models) > 0 { + m.store.Cache.Set(key, models) + } + return models, nil +} + +func parseModelsForProtocol(protocol string, body []byte) []ModelEntry { + if protocol == "replicate" { + if models := parseReplicateModels(body); len(models) > 0 { + return models + } + } + if models := parseOpenAIModels(body); len(models) > 0 { + return models + } + if protocol == "replicate" { + return parseReplicateModels(body) + } + return nil +} + +func parseOpenAIModels(body []byte) []ModelEntry { + var openai struct { + Data []struct { + ID string `json:"id"` + Name string `json:"name"` + } `json:"data"` + } + if err := json.Unmarshal(body, &openai); err == nil && len(openai.Data) > 0 { + out := make([]ModelEntry, 0, len(openai.Data)) + for _, item := range openai.Data { + name := item.Name + if name == "" { + name = item.ID + } + if item.ID != "" { + out = append(out, ModelEntry{ID: item.ID, Name: name}) + } + } + return out + } + var arr []struct { + ID string `json:"id"` + Name string `json:"name"` + } + if err := json.Unmarshal(body, &arr); err == nil && len(arr) > 0 { + out := make([]ModelEntry, 0, len(arr)) + for _, item := range arr { + name := item.Name + if name == "" { + name = item.ID + } + if item.ID != "" { + out = append(out, ModelEntry{ID: item.ID, Name: name}) + } + } + return out + } + return nil +} + +func parseReplicateModels(body []byte) []ModelEntry { + var resp struct { + Results []struct { + Owner string `json:"owner"` + Name string `json:"name"` + Description string `json:"description"` + } `json:"results"` + } + if err := json.Unmarshal(body, &resp); err != nil || len(resp.Results) == 0 { + return nil + } + out := make([]ModelEntry, 0, len(resp.Results)) + for _, item := range resp.Results { + if item.Owner == "" || item.Name == "" { + continue + } + id := item.Owner + "/" + item.Name + name := item.Name + if item.Description != "" { + name = item.Description + } else { + name = id + } + out = append(out, ModelEntry{ID: id, Name: name}) + } + return out +} + +func parseDetailsForProtocol(protocol string, body []byte) []ModelDetail { + if protocol == "replicate" { + if models := parseReplicateModelDetails(body); len(models) > 0 { + return models + } + } + if entries := parseOpenAIModels(body); len(entries) > 0 { + out := make([]ModelDetail, len(entries)) + for i, e := range entries { + out[i] = manualModelDetail(e.ID, e.Name) + } + return out + } + if protocol == "replicate" { + if entries := parseReplicateModels(body); len(entries) > 0 { + out := make([]ModelDetail, len(entries)) + for i, e := range entries { + out[i] = manualModelDetail(e.ID, e.Name) + } + return out + } + } + return nil +} + +func mergeModelDetails(upstream []ModelDetail, manual []ModelEntry) []ModelDetail { + seen := map[string]struct{}{} + out := make([]ModelDetail, 0, len(upstream)+len(manual)) + for _, m := range upstream { + if _, ok := seen[m.ID]; ok { + continue + } + seen[m.ID] = struct{}{} + out = append(out, m) + } + for _, m := range manual { + if _, ok := seen[m.ID]; ok { + continue + } + seen[m.ID] = struct{}{} + out = append(out, manualModelDetail(m.ID, m.Name)) + } + return out +} + +func mergeModels(upstream, manual []ModelEntry) []ModelEntry { + seen := map[string]struct{}{} + out := make([]ModelEntry, 0, len(upstream)+len(manual)) + for _, m := range upstream { + if _, ok := seen[m.ID]; ok { + continue + } + seen[m.ID] = struct{}{} + out = append(out, m) + } + for _, m := range manual { + if _, ok := seen[m.ID]; ok { + continue + } + seen[m.ID] = struct{}{} + out = append(out, m) + } + return out +} + +func cacheKey(providerID int64) string { + return fmt.Sprintf("upstream_models:%d", providerID) +} + +func detailsCacheKey(providerID int64) string { + return fmt.Sprintf("upstream_model_details:%d", providerID) +} + +func (m *ModelsService) invalidateCache(providerID int64) { + m.store.Cache.Delete(cacheKey(providerID)) + m.store.Cache.Delete(detailsCacheKey(providerID)) +} + +func FilterModelsByKind(models []ModelDetail, kind string) []ModelDetail { + out := make([]ModelDetail, 0, len(models)) + for _, m := range models { + if m.Kind == kind { + out = append(out, m) + } + } + return out +} + +func DefaultModelDetail(models []ModelDetail) string { + if len(models) > 0 { + return models[0].ID + } + return "" +} diff --git a/internal/provider/models_test.go b/internal/provider/models_test.go new file mode 100644 index 0000000..f48cb4b --- /dev/null +++ b/internal/provider/models_test.go @@ -0,0 +1,44 @@ +package provider + +import "testing" + +func TestModelsPath(t *testing.T) { + openai := ModelsPath(&Config{Protocol: "openai", ServiceType: "chat"}) + if openai != "/v1/models?protocol=openai&service_type=chat" { + t.Fatalf("openai path: %s", openai) + } + replicate := ModelsPath(&Config{Protocol: "replicate", ServiceType: "text_to_image"}) + if replicate != "/v1/models?protocol=replicate&service_type=text_to_image" { + t.Fatalf("replicate path: %s", replicate) + } +} + +func TestParseReplicateModels(t *testing.T) { + body := []byte(`{ + "results": [ + {"owner":"stability-ai","name":"sdxl","description":"SDXL"}, + {"owner":"black-forest-labs","name":"flux-schnell"} + ] + }`) + models := parseReplicateModels(body) + if len(models) != 2 { + t.Fatalf("expected 2 models, got %d", len(models)) + } + if models[0].ID != "stability-ai/sdxl" || models[0].Name != "SDXL" { + t.Fatalf("unexpected first model: %+v", models[0]) + } + if models[1].ID != "black-forest-labs/flux-schnell" { + t.Fatalf("unexpected second model id: %s", models[1].ID) + } +} + +func TestParseModelsForProtocol(t *testing.T) { + openaiBody := []byte(`{"data":[{"id":"gpt-4o","name":"GPT-4o"}]}`) + if models := parseModelsForProtocol("openai", openaiBody); len(models) != 1 || models[0].ID != "gpt-4o" { + t.Fatalf("openai parse failed: %+v", models) + } + replicateBody := []byte(`{"results":[{"owner":"a","name":"b"}]}`) + if models := parseModelsForProtocol("replicate", replicateBody); len(models) != 1 || models[0].ID != "a/b" { + t.Fatalf("replicate parse failed: %+v", models) + } +} diff --git a/internal/provider/router.go b/internal/provider/router.go new file mode 100644 index 0000000..1180293 --- /dev/null +++ b/internal/provider/router.go @@ -0,0 +1,95 @@ +package provider + +import ( + "context" + "fmt" + "net/http" + "strings" + "time" + + "github.com/rose_cat707/Atlas/internal/service/reqlog" + "github.com/rose_cat707/Atlas/internal/store" +) + +type Config struct { + ID int64 + Name string + Protocol string + BaseURL string + APIKey string + ServiceType string + ExtraJSON string +} + +type Router struct { + store *store.Store + client *http.Client + streamClient *http.Client + logger *reqlog.Logger +} + +func NewRouter(st *store.Store, logger *reqlog.Logger) *Router { + return &Router{ + store: st, + client: &http.Client{ + Timeout: 120 * time.Second, + }, + streamClient: &http.Client{}, + logger: logger, + } +} + +func (r *Router) Resolve(ctx context.Context, serviceType string, providerID int64) (*Config, error) { + var p *store.Provider + var err error + if providerID > 0 { + p, err = r.store.GetProvider(ctx, providerID) + } else { + p, err = r.store.DefaultProvider(ctx, serviceType) + } + if err != nil { + return nil, fmt.Errorf("no provider for %s: %w", serviceType, err) + } + return r.providerToConfig(p) +} + +func (r *Router) ResolveByProtocol(ctx context.Context, serviceType, protocol string, providerID int64) (*Config, error) { + if providerID > 0 { + p, err := r.store.GetProvider(ctx, providerID) + if err != nil { + return nil, err + } + return r.providerToConfig(p) + } + providers, err := r.store.ListProviders(ctx, serviceType, true) + if err != nil { + return nil, err + } + for _, p := range providers { + if protocol == "" || p.Protocol == protocol { + return r.providerToConfig(&p) + } + } + return nil, fmt.Errorf("no provider for %s/%s", serviceType, protocol) +} + +func (r *Router) providerToConfig(p *store.Provider) (*Config, error) { + if !p.Enabled { + return nil, fmt.Errorf("provider %d disabled", p.ID) + } + return &Config{ + ID: p.ID, + Name: p.Name, + Protocol: p.Protocol, + BaseURL: strings.TrimRight(p.BaseURL, "/"), + APIKey: p.APIKey, + ServiceType: p.ServiceType, + ExtraJSON: p.ExtraJSON, + }, nil +} + +type ProxyResult struct { + StatusCode int + Body []byte + Headers http.Header +} diff --git a/internal/provider/stream.go b/internal/provider/stream.go new file mode 100644 index 0000000..7761be4 --- /dev/null +++ b/internal/provider/stream.go @@ -0,0 +1,143 @@ +package provider + +import ( + "bufio" + "bytes" + "context" + "fmt" + "io" + "net/http" + "strings" + "time" + + "github.com/rose_cat707/Atlas/internal/service/reqlog" +) + +type CallMeta struct { + SessionID string + Method string + Path string + Model string +} + +func (r *Router) ProxyWithMeta(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string, meta CallMeta) (*ProxyResult, error) { + start := time.Now() + res, err := r.proxy(ctx, cfg, method, path, body, headers) + r.record(meta, cfg, start, res, err, body, "") + return res, err +} + +func (r *Router) ProxyStream(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string, meta CallMeta, w io.Writer, flush func()) error { + start := time.Now() + url := cfg.BaseURL + path + req, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body)) + if err != nil { + r.record(meta, cfg, start, nil, err, body, "") + return err + } + req.Header.Set("Authorization", "Bearer "+cfg.APIKey) + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + for k, v := range headers { + req.Header.Set(k, v) + } + + resp, err := r.streamClient.Do(req) + if err != nil { + r.record(meta, cfg, start, nil, err, body, "") + return err + } + defer resp.Body.Close() + + if resp.StatusCode >= 400 { + data, _ := io.ReadAll(resp.Body) + res := &ProxyResult{StatusCode: resp.StatusCode, Body: data, Headers: resp.Header} + r.record(meta, cfg, start, res, fmt.Errorf("upstream %d", resp.StatusCode), body, "") + return fmt.Errorf("upstream %d: %s", resp.StatusCode, string(data)) + } + + var capture strings.Builder + reader := bufio.NewReader(resp.Body) + for { + line, readErr := reader.ReadString('\n') + if len(line) > 0 { + capture.WriteString(line) + if _, err := io.WriteString(w, line); err != nil { + return err + } + if flush != nil && (line == "\n" || strings.HasPrefix(line, "data:")) { + flush() + } + } + if readErr == io.EOF { + break + } + if readErr != nil { + r.record(meta, cfg, start, &ProxyResult{StatusCode: resp.StatusCode}, readErr, body, capture.String()) + return readErr + } + } + r.record(meta, cfg, start, &ProxyResult{StatusCode: resp.StatusCode}, nil, body, capture.String()) + return nil +} + +func (r *Router) record(meta CallMeta, cfg *Config, start time.Time, res *ProxyResult, callErr error, reqBody []byte, streamCapture string) { + if r.logger == nil { + return + } + entry := reqlog.Entry{ + SessionID: meta.SessionID, + Method: meta.Method, + Path: meta.Path, + Protocol: cfg.Protocol, + ProviderID: reqlog.ProviderIDPtr(cfg.ID), + Model: meta.Model, + LatencyMS: reqlog.Since(start), + RequestSummary: reqlog.Summarize(reqBody, 512), + } + if res != nil { + entry.StatusCode = res.StatusCode + if streamCapture != "" { + entry.ResponseSummary = reqlog.SummarizeSSE(streamCapture, 512) + } else { + entry.ResponseSummary = reqlog.Summarize(res.Body, 512) + } + } + if callErr != nil { + entry.Error = callErr.Error() + if entry.StatusCode == 0 { + entry.StatusCode = 502 + } + } + r.logger.Log(context.Background(), entry) +} + +func (r *Router) proxy(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string) (*ProxyResult, error) { + url := cfg.BaseURL + path + req, err := http.NewRequestWithContext(ctx, method, url, bytes.NewReader(body)) + if err != nil { + return nil, err + } + req.Header.Set("Authorization", "Bearer "+cfg.APIKey) + if len(body) > 0 { + req.Header.Set("Content-Type", "application/json") + } + for k, v := range headers { + req.Header.Set(k, v) + } + resp, err := r.client.Do(req) + if err != nil { + return nil, err + } + defer resp.Body.Close() + data, err := io.ReadAll(resp.Body) + if err != nil { + return nil, err + } + return &ProxyResult{StatusCode: resp.StatusCode, Body: data, Headers: resp.Header}, nil +} + +// Proxy keeps backward compatibility without logging. +func (r *Router) Proxy(ctx context.Context, cfg *Config, method, path string, body []byte, headers map[string]string) (*ProxyResult, error) { + return r.proxy(ctx, cfg, method, path, body, headers) +} diff --git a/internal/provider/testdata/models_sample.json b/internal/provider/testdata/models_sample.json new file mode 100644 index 0000000..7d10c5f --- /dev/null +++ b/internal/provider/testdata/models_sample.json @@ -0,0 +1,75 @@ +{ + "results": [ + { + "description": "Qwen Image Edit", + "latest_version": { + "id": "qwen-image-edit", + "openapi_schema": { + "components": { + "schemas": { + "Input": { + "properties": { + "prompt": { + "description": "正向提示词", + "title": "正向提示词", + "type": "string", + "x-order": 0 + }, + "image": { + "description": "输入图 URL", + "title": "输入图 URL", + "type": "string", + "x-order": 2 + }, + "denoise": { + "default": 0.75, + "description": "重绘幅度", + "title": "重绘幅度", + "type": "number", + "x-order": 6 + } + }, + "required": ["prompt", "image"], + "type": "object" + } + } + } + } + }, + "name": "qwen-image-edit", + "owner": "comfyui" + }, + { + "description": "Z Image", + "latest_version": { + "id": "z-image", + "openapi_schema": { + "components": { + "schemas": { + "Input": { + "properties": { + "prompt": { + "description": "正向提示词", + "title": "正向提示词", + "type": "string", + "x-order": 0 + }, + "width": { + "default": 1024, + "title": "宽度", + "type": "integer", + "x-order": 3 + } + }, + "required": ["prompt"], + "type": "object" + } + } + } + } + }, + "name": "z-image", + "owner": "comfyui" + } + ] +} diff --git a/internal/service/app.go b/internal/service/app.go index d1a1101..4b9d74d 100644 --- a/internal/service/app.go +++ b/internal/service/app.go @@ -1,8 +1,41 @@ package service -import "github.com/rose_cat707/go-vue-template/internal/store" +import ( + "github.com/rose_cat707/Atlas/internal/provider" + "github.com/rose_cat707/Atlas/internal/service/asset" + "github.com/rose_cat707/Atlas/internal/service/chat" + "github.com/rose_cat707/Atlas/internal/service/invoke" + "github.com/rose_cat707/Atlas/internal/service/reqlog" + "github.com/rose_cat707/Atlas/internal/service/taskevents" + "github.com/rose_cat707/Atlas/internal/store" +) type App struct { - Store *store.Store - DataDir string + Store *store.Store + DataDir string + Router *provider.Router + Models *provider.ModelsService + TaskEvents *taskevents.Hub + Invoke *invoke.Service + Chat *chat.Service + Assets *asset.Signer + ReqLog *reqlog.Logger +} + +func NewApp(st *store.Store, dataDir string, signSecret string) *App { + logger := reqlog.New(st) + router := provider.NewRouter(st, logger) + models := provider.NewModelsService(st, router) + taskEvents := taskevents.NewHub() + return &App{ + Store: st, + DataDir: dataDir, + Router: router, + Models: models, + TaskEvents: taskEvents, + Invoke: invoke.NewService(st, router, models, taskEvents), + Chat: chat.NewService(st, router, models), + Assets: asset.NewSigner(signSecret), + ReqLog: logger, + } } diff --git a/internal/service/asset/sign.go b/internal/service/asset/sign.go new file mode 100644 index 0000000..dd9f3a6 --- /dev/null +++ b/internal/service/asset/sign.go @@ -0,0 +1,85 @@ +package asset + +import ( + "crypto/hmac" + "crypto/sha256" + "encoding/hex" + "fmt" + "io" + "mime" + "os" + "path/filepath" + "strconv" + "strings" + "time" + + "github.com/google/uuid" +) + +type Signer struct { + secret string +} + +func NewSigner(secret string) *Signer { + return &Signer{secret: secret} +} + +func (s *Signer) Sign(fileID int64, exp int64) string { + mac := hmac.New(sha256.New, []byte(s.secret)) + mac.Write([]byte(fmt.Sprintf("%d:%d", fileID, exp))) + return hex.EncodeToString(mac.Sum(nil)) +} + +func (s *Signer) Verify(fileID int64, exp int64, sig string) bool { + if time.Now().Unix() > exp { + return false + } + expected := s.Sign(fileID, exp) + return hmac.Equal([]byte(expected), []byte(sig)) +} + +func (s *Signer) SignedURL(basePath string, fileID int64, ttl time.Duration) string { + exp := time.Now().Add(ttl).Unix() + sig := s.Sign(fileID, exp) + return fmt.Sprintf("%s/user/files/%d?exp=%d&sig=%s", basePath, fileID, exp, sig) +} + +func SaveUpload(dataDir, sessionID string, filename string, r io.Reader, maxBytes int64) (string, int64, string, error) { + ext := filepath.Ext(filename) + if ext == "" { + ext = ".bin" + } + id := uuid.New().String() + dir := filepath.Join(dataDir, "files", sessionID, "uploads") + if err := os.MkdirAll(dir, 0o755); err != nil { + return "", 0, "", err + } + path := filepath.Join(dir, id+ext) + f, err := os.Create(path) + if err != nil { + return "", 0, "", err + } + defer f.Close() + n, err := io.Copy(f, io.LimitReader(r, maxBytes+1)) + if err != nil { + os.Remove(path) + return "", 0, "", err + } + if n > maxBytes { + os.Remove(path) + return "", 0, "", fmt.Errorf("file too large") + } + mimeType := mime.TypeByExtension(ext) + if mimeType == "" { + mimeType = "application/octet-stream" + } + return path, n, mimeType, nil +} + +func ParseSignedQuery(expStr, sig string) (int64, string, error) { + exp, err := strconv.ParseInt(expStr, 10, 64) + if err != nil { + return 0, "", err + } + return exp, strings.TrimSpace(sig), nil +} diff --git a/internal/service/chat/service.go b/internal/service/chat/service.go new file mode 100644 index 0000000..e3c0d88 --- /dev/null +++ b/internal/service/chat/service.go @@ -0,0 +1,210 @@ +package chat + +import ( + "context" + "encoding/json" + "fmt" + "io" + "strings" + + "github.com/rose_cat707/Atlas/internal/domain" + "github.com/rose_cat707/Atlas/internal/provider" + "github.com/rose_cat707/Atlas/internal/store" +) + +type Service struct { + store *store.Store + router *provider.Router + models *provider.ModelsService +} + +func NewService(st *store.Store, router *provider.Router, models *provider.ModelsService) *Service { + return &Service{store: st, router: router, models: models} +} + +type StreamRequest struct { + SessionID string + ConversationID int64 + Content string + Model string + ProviderID int64 +} + +type streamCapture struct { + io.Writer + buf []byte +} + +func (s *streamCapture) Write(p []byte) (int, error) { + s.buf = append(s.buf, p...) + return s.Writer.Write(p) +} + +func (s *Service) StreamMessage(ctx context.Context, req StreamRequest, w io.Writer, flusher func()) error { + conv, err := s.store.GetConversation(ctx, req.SessionID, req.ConversationID) + if err != nil { + return err + } + userMsg := &store.ChatMessageRow{ConversationID: req.ConversationID, Role: "user", Content: req.Content} + if _, err := s.store.InsertMessage(ctx, userMsg); err != nil { + return err + } + msgs, err := s.store.ListMessages(ctx, req.ConversationID, 40) + if err != nil { + return err + } + model := req.Model + if model == "" { + model = conv.Model + } + cfg, err := s.router.Resolve(ctx, string(domain.ServiceChat), req.ProviderID) + if err != nil { + return err + } + if model == "" { + models, _ := s.models.ListEffective(ctx, cfg) + model = provider.DefaultModel(models) + } + body, err := json.Marshal(map[string]any{ + "model": model, + "messages": toChatMessages(msgs), + "stream": true, + }) + if err != nil { + return err + } + capture := &streamCapture{Writer: w} + meta := provider.CallMeta{ + SessionID: req.SessionID, + Method: "POST", + Path: "/v1/chat/completions", + Model: model, + } + if err := s.router.ProxyStream(ctx, cfg, "POST", "/v1/chat/completions", body, nil, meta, capture, flusher); err != nil { + return err + } + assistant := extractAssistantFromSSE(string(capture.buf)) + return s.persistAssistant(ctx, req, model, cfg, assistant) +} + +func (s *Service) persistAssistant(ctx context.Context, req StreamRequest, model string, cfg *provider.Config, assistant string) error { + asst := &store.ChatMessageRow{ConversationID: req.ConversationID, Role: "assistant", Content: assistant} + if _, err := s.store.InsertMessage(ctx, asst); err != nil { + return err + } + _ = s.maybeSummarizeTitle(ctx, req, model, cfg) + + pid := cfg.ID + out, _ := json.Marshal(map[string]any{"content": assistant}) + gen := &store.Generation{ + SessionID: req.SessionID, + ServiceType: string(domain.ServiceChat), + ProviderID: &pid, + Model: model, + InputJSON: fmt.Sprintf(`{"content":%q}`, req.Content), + OutputJSON: string(out), + } + s.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error { + _, err := st.InsertGeneration(ctx, gen) + return err + }) + return nil +} + +func (s *Service) maybeSummarizeTitle(ctx context.Context, req StreamRequest, model string, cfg *provider.Config) error { + conv, err := s.store.GetConversation(ctx, req.SessionID, req.ConversationID) + if err != nil { + return err + } + params := conversationParams(conv.ParamsJSON) + if titleAlreadySummarized(params) { + return nil + } + userCount, err := s.store.CountUserMessages(ctx, req.ConversationID) + if err != nil { + return err + } + if userCount <= titleSummaryRoundThreshold { + return nil + } + + msgs, err := s.store.ListMessages(ctx, req.ConversationID, 24) + if err != nil { + return err + } + title, err := s.generateTitle(ctx, req, model, cfg, msgs) + if err != nil || title == "" { + return err + } + conv.Title = title + if model != "" { + conv.Model = model + } + params["title_summarized"] = true + conv.ParamsJSON = conversationParamsJSON(params) + return s.store.UpdateConversation(ctx, conv) +} + +func (s *Service) generateTitle(ctx context.Context, req StreamRequest, model string, cfg *provider.Config, msgs []store.ChatMessageRow) (string, error) { + chatMsgs := []map[string]string{{"role": "system", "content": titleSummarySystemPrompt}} + chatMsgs = append(chatMsgs, toChatMessages(msgs)...) + body, err := json.Marshal(map[string]any{ + "model": model, + "messages": chatMsgs, + "max_tokens": 40, + "temperature": 0.3, + "stream": false, + }) + if err != nil { + return "", err + } + meta := provider.CallMeta{ + SessionID: req.SessionID, + Method: "POST", + Path: "/v1/chat/completions", + Model: model, + } + res, err := s.router.ProxyWithMeta(ctx, cfg, "POST", "/v1/chat/completions", body, nil, meta) + if err != nil { + return "", err + } + if res.StatusCode >= 400 { + return "", fmt.Errorf("title summarize upstream %d", res.StatusCode) + } + return sanitizeTitle(parseChatCompletionContent(res.Body)), nil +} + +func extractAssistantFromSSE(raw string) string { + var assistant strings.Builder + for _, line := range strings.Split(raw, "\n") { + if !strings.HasPrefix(line, "data: ") { + continue + } + data := strings.TrimPrefix(line, "data: ") + if data == "[DONE]" { + break + } + var chunk struct { + Choices []struct { + Delta struct { + Content string `json:"content"` + } `json:"delta"` + } `json:"choices"` + } + if err := json.Unmarshal([]byte(data), &chunk); err != nil { + continue + } + if len(chunk.Choices) > 0 { + assistant.WriteString(chunk.Choices[0].Delta.Content) + } + } + return assistant.String() +} + +func toChatMessages(rows []store.ChatMessageRow) []map[string]string { + out := make([]map[string]string, 0, len(rows)) + for _, m := range rows { + out = append(out, map[string]string{"role": m.Role, "content": m.Content}) + } + return out +} diff --git a/internal/service/chat/title.go b/internal/service/chat/title.go new file mode 100644 index 0000000..2af6e05 --- /dev/null +++ b/internal/service/chat/title.go @@ -0,0 +1,68 @@ +package chat + +import ( + "encoding/json" + "strings" + "unicode/utf8" +) + +const ( + titleSummaryRoundThreshold = 3 + titleSummarySystemPrompt = "你是会话标题生成器。根据对话内容输出一条简洁的中文标题，不超过20字，不要引号、书名号或句号，只输出标题本身。" +) + +func conversationParams(raw string) map[string]any { + out := map[string]any{} + if raw == "" { + return out + } + _ = json.Unmarshal([]byte(raw), &out) + return out +} + +func conversationParamsJSON(params map[string]any) string { + if len(params) == 0 { + return "{}" + } + data, err := json.Marshal(params) + if err != nil { + return "{}" + } + return string(data) +} + +func titleAlreadySummarized(params map[string]any) bool { + v, ok := params["title_summarized"].(bool) + return ok && v +} + +func sanitizeTitle(title string) string { + title = strings.TrimSpace(title) + title = strings.Trim(title, `"'""''《》`) + title = strings.ReplaceAll(title, "\n", " ") + if title == "" { + return "" + } + if utf8.RuneCountInString(title) > 30 { + runes := []rune(title) + title = string(runes[:30]) + } + return title +} + +func parseChatCompletionContent(body []byte) string { + var resp struct { + Choices []struct { + Message struct { + Content string `json:"content"` + } `json:"message"` + } `json:"choices"` + } + if err := json.Unmarshal(body, &resp); err != nil { + return "" + } + if len(resp.Choices) == 0 { + return "" + } + return strings.TrimSpace(resp.Choices[0].Message.Content) +} diff --git a/internal/service/chat/title_test.go b/internal/service/chat/title_test.go new file mode 100644 index 0000000..7ab58d5 --- /dev/null +++ b/internal/service/chat/title_test.go @@ -0,0 +1,38 @@ +package chat + +import "testing" + +func TestSanitizeTitle(t *testing.T) { + if got := sanitizeTitle(` "测试标题" `); got != "测试标题" { + t.Fatalf("unexpected: %q", got) + } + long := stringsRepeat("标题", 20) + got := sanitizeTitle(long) + if len([]rune(got)) > 30 { + t.Fatalf("title too long: %d", len([]rune(got))) + } +} + +func TestTitleAlreadySummarized(t *testing.T) { + if titleAlreadySummarized(map[string]any{}) { + t.Fatal("expected false") + } + if !titleAlreadySummarized(map[string]any{"title_summarized": true}) { + t.Fatal("expected true") + } +} + +func stringsRepeat(s string, n int) string { + out := "" + for i := 0; i < n; i++ { + out += s + } + return out +} + +func TestParseChatCompletionContent(t *testing.T) { + body := []byte(`{"choices":[{"message":{"content":" Atlas 部署方案 "}}]}`) + if got := parseChatCompletionContent(body); got != "Atlas 部署方案" { + t.Fatalf("unexpected: %q", got) + } +} diff --git a/internal/service/invoke/async_worker.go b/internal/service/invoke/async_worker.go new file mode 100644 index 0000000..eddf0fa --- /dev/null +++ b/internal/service/invoke/async_worker.go @@ -0,0 +1,140 @@ +package invoke + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/rose_cat707/Atlas/internal/provider" + "github.com/rose_cat707/Atlas/internal/store" +) + +type asyncWorker struct { + svc *Service +} + +func newAsyncWorker(svc *Service) *asyncWorker { + return &asyncWorker{svc: svc} +} + +func (w *asyncWorker) run() { + ticker := time.NewTicker(2 * time.Second) + defer ticker.Stop() + for range ticker.C { + w.pollOnce() + } +} + +func (w *asyncWorker) pollOnce() { + ctx := context.Background() + tasks, err := w.svc.store.ListPendingTasks(ctx) + if err != nil { + return + } + for _, task := range tasks { + w.pollTask(ctx, task) + } +} + +func (w *asyncWorker) pollTask(ctx context.Context, task store.Task) { + cfg, err := w.svc.store.GetProvider(ctx, task.ProviderID) + if err != nil { + return + } + pcfg := &provider.Config{ + ID: cfg.ID, + Protocol: cfg.Protocol, + BaseURL: cfg.BaseURL, + APIKey: cfg.APIKey, + } + path := fmt.Sprintf("/v1/predictions/%s", task.ExternalID) + res, err := w.svc.router.Proxy(ctx, pcfg, "GET", path, nil, nil) + if err != nil { + return + } + if res.StatusCode >= 400 { + task.Status = "failed" + task.Error = string(res.Body) + w.svc.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error { + return st.UpdateTask(ctx, &task) + }) + w.svc.notifyTask(task) + return + } + var pred struct { + Status string `json:"status"` + Output any `json:"output"` + Error string `json:"error"` + Progress int `json:"progress"` + } + _ = json.Unmarshal(res.Body, &pred) + task.Progress = pred.Progress + switch pred.Status { + case "succeeded": + task.Status = "succeeded" + now := time.Now().UTC().Format(time.RFC3339) + task.CompletedAt = &now + out, _ := json.Marshal(map[string]any{"output": pred.Output}) + task.ResultJSON = string(out) + pid := task.ProviderID + gen := &store.Generation{ + SessionID: task.SessionID, + ServiceType: task.ServiceType, + ProviderID: &pid, + InputJSON: task.InputJSON, + OutputJSON: task.ResultJSON, + TaskID: &task.ID, + } + w.svc.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error { + _, err := st.InsertGeneration(ctx, gen) + return err + }) + case "failed", "canceled": + task.Status = pred.Status + task.Error = pred.Error + now := time.Now().UTC().Format(time.RFC3339) + task.CompletedAt = &now + case "processing", "starting": + task.Status = "running" + default: + task.Status = "running" + } + w.svc.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error { + return st.UpdateTask(ctx, &task) + }) + w.svc.notifyTask(task) +} + +func (s *Service) CancelTask(ctx context.Context, sessionID string, taskID int64) error { + task, err := s.store.GetTask(ctx, sessionID, taskID) + if err != nil { + return err + } + cfg, err := s.store.GetProvider(ctx, task.ProviderID) + if err != nil { + return err + } + pcfg := &provider.Config{BaseURL: cfg.BaseURL, APIKey: cfg.APIKey} + path := fmt.Sprintf("/v1/predictions/%s/cancel", task.ExternalID) + meta := provider.CallMeta{ + SessionID: sessionID, + Method: "POST", + Path: "/v1/predictions/:id/cancel", + } + res, err := s.router.ProxyWithMeta(ctx, pcfg, "POST", path, nil, nil, meta) + if err != nil { + return err + } + if res.StatusCode >= 400 { + return fmt.Errorf("cancel failed: %s", string(res.Body)) + } + task.Status = "canceled" + now := time.Now().UTC().Format(time.RFC3339) + task.CompletedAt = &now + if err := s.store.UpdateTask(ctx, task); err != nil { + return err + } + s.notifyTask(*task) + return nil +} diff --git a/internal/service/invoke/service.go b/internal/service/invoke/service.go new file mode 100644 index 0000000..4f4a833 --- /dev/null +++ b/internal/service/invoke/service.go @@ -0,0 +1,251 @@ +package invoke + +import ( + "context" + "encoding/json" + "fmt" + "time" + + "github.com/rose_cat707/Atlas/internal/domain" + "github.com/rose_cat707/Atlas/internal/provider" + "github.com/rose_cat707/Atlas/internal/service/taskevents" + "github.com/rose_cat707/Atlas/internal/store" +) + +type Service struct { + store *store.Store + router *provider.Router + models *provider.ModelsService + taskEvents *taskevents.Hub + specs map[domain.ServiceType]*domain.ServiceTypeSpec + worker *asyncWorker +} + +func NewService(st *store.Store, router *provider.Router, models *provider.ModelsService, taskEvents *taskevents.Hub) *Service { + s := &Service{ + store: st, + router: router, + models: models, + taskEvents: taskEvents, + specs: defaultSpecs(), + } + s.worker = newAsyncWorker(s) + go s.worker.run() + return s +} + +func (s *Service) Spec(t domain.ServiceType) (*domain.ServiceTypeSpec, bool) { + spec, ok := s.specs[t] + return spec, ok +} + +func (s *Service) Specs() map[domain.ServiceType]*domain.ServiceTypeSpec { + return s.specs +} + +func (s *Service) Run(ctx context.Context, req domain.InvokeRequest) (*domain.InvokeResult, error) { + spec, ok := s.specs[domain.ServiceType(req.ServiceType)] + if !ok { + return nil, fmt.Errorf("unknown service_type: %s", req.ServiceType) + } + switch req.Mode { + case domain.ModeStream: + if spec.Capabilities.Stream == nil { + return nil, fmt.Errorf("stream not supported for %s", req.ServiceType) + } + case domain.ModeSync: + if spec.Capabilities.Sync == nil { + return nil, fmt.Errorf("sync not supported for %s", req.ServiceType) + } + case domain.ModeAsync: + if spec.Capabilities.Async == nil { + return nil, fmt.Errorf("async not supported for %s", req.ServiceType) + } + default: + return nil, fmt.Errorf("invalid mode: %s", req.Mode) + } + + var params map[string]any + if len(req.Params) > 0 { + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + return nil, fmt.Errorf("invalid params: %w", err) + } + } + if spec.ValidateParams != nil { + if err := spec.ValidateParams(params); err != nil { + return nil, err + } + } + + cfg, err := s.router.Resolve(ctx, string(req.ServiceType), req.ProviderID) + if err != nil { + return nil, err + } + if req.Model == "" { + models, _ := s.models.ListEffective(ctx, cfg) + req.Model = provider.DefaultModel(models) + } + + if isImageService(req.ServiceType) { + detail, err := s.models.FindModelDetail(ctx, cfg, req.Model) + if err != nil { + return nil, err + } + if err := provider.ValidateModelInput(detail.InputSchema, params); err != nil { + return nil, err + } + } + + switch req.Mode { + case domain.ModeAsync: + return s.runAsync(ctx, req, cfg, params) + case domain.ModeSync: + return s.runSync(ctx, req, cfg, params) + default: + return nil, fmt.Errorf("stream mode handled by chat service") + } +} + +func (s *Service) runSync(ctx context.Context, req domain.InvokeRequest, cfg *provider.Config, params map[string]any) (*domain.InvokeResult, error) { + body, err := buildPredictionBody(req.Model, params) + if err != nil { + return nil, err + } + meta := provider.CallMeta{ + SessionID: req.SessionID, + Method: "POST", + Path: "/v1/predictions", + Model: req.Model, + } + res, err := s.router.ProxyWithMeta(ctx, cfg, "POST", "/v1/predictions", body, map[string]string{"Prefer": "wait"}, meta) + if err != nil { + return nil, err + } + if res.StatusCode >= 400 { + return nil, fmt.Errorf("upstream error %d: %s", res.StatusCode, string(res.Body)) + } + output := parsePredictionOutput(res.Body) + pid := cfg.ID + gen := &store.Generation{ + SessionID: req.SessionID, + ServiceType: string(req.ServiceType), + ProviderID: &pid, + Model: req.Model, + InputJSON: string(req.Params), + OutputJSON: string(output), + } + id, err := s.store.InsertGeneration(ctx, gen) + if err != nil { + return nil, err + } + return &domain.InvokeResult{GenerationID: id, Output: output}, nil +} + +func (s *Service) runAsync(ctx context.Context, req domain.InvokeRequest, cfg *provider.Config, params map[string]any) (*domain.InvokeResult, error) { + body, err := buildPredictionBody(req.Model, params) + if err != nil { + return nil, err + } + meta := provider.CallMeta{ + SessionID: req.SessionID, + Method: "POST", + Path: "/v1/predictions", + Model: req.Model, + } + res, err := s.router.ProxyWithMeta(ctx, cfg, "POST", "/v1/predictions", body, nil, meta) + if err != nil { + return nil, err + } + if res.StatusCode >= 400 { + return nil, fmt.Errorf("upstream error %d: %s", res.StatusCode, string(res.Body)) + } + var pred struct { + ID string `json:"id"` + } + _ = json.Unmarshal(res.Body, &pred) + task := &store.Task{ + SessionID: req.SessionID, + ServiceType: string(req.ServiceType), + Protocol: cfg.Protocol, + ProviderID: cfg.ID, + Status: "running", + Progress: 0, + InputJSON: string(req.Params), + ExternalID: pred.ID, + } + taskID, err := s.store.InsertTask(ctx, task) + if err != nil { + return nil, err + } + task.ID = taskID + s.notifyTask(*task) + return &domain.InvokeResult{TaskID: &taskID}, nil +} + +func (s *Service) notifyTask(task store.Task) { + s.store.Cache.Set(fmt.Sprintf("task:%d", task.ID), &task) + if s.taskEvents != nil { + s.taskEvents.Publish(task) + } +} + +func buildPredictionBody(model string, params map[string]any) ([]byte, error) { + input := map[string]any{} + for k, v := range params { + input[k] = v + } + payload := map[string]any{ + "model": model, + "input": input, + } + return json.Marshal(payload) +} + +func isImageService(serviceType domain.ServiceType) bool { + return serviceType == domain.ServiceTextToImage || serviceType == domain.ServiceImageToImage +} + +func parsePredictionOutput(body []byte) json.RawMessage { + var pred map[string]any + if err := json.Unmarshal(body, &pred); err != nil { + return json.RawMessage(`{}`) + } + out := map[string]any{} + if output, ok := pred["output"]; ok { + out["output"] = output + } + if urls, ok := pred["urls"]; ok { + out["urls"] = urls + } + data, _ := json.Marshal(out) + return data +} + +func defaultSpecs() map[domain.ServiceType]*domain.ServiceTypeSpec { + return map[domain.ServiceType]*domain.ServiceTypeSpec{ + domain.ServiceChat: { + Type: domain.ServiceChat, + Capabilities: domain.InvokeCapabilities{ + Stream: &domain.StreamCapability{}, + }, + DefaultParams: func() map[string]any { return map[string]any{"temperature": 0.7} }, + ValidateParams: func(params map[string]any) error { return nil }, + }, + domain.ServiceTextToImage: { + Type: domain.ServiceTextToImage, + Capabilities: domain.InvokeCapabilities{ + Sync: &domain.SyncCapability{}, + Async: &domain.AsyncCapability{Driver: "replicate", PollInterval: 2 * time.Second}, + }, + DefaultParams: func() map[string]any { return map[string]any{} }, + }, + domain.ServiceImageToImage: { + Type: domain.ServiceImageToImage, + Capabilities: domain.InvokeCapabilities{ + Sync: &domain.SyncCapability{}, + Async: &domain.AsyncCapability{Driver: "replicate", PollInterval: 2 * time.Second}, + }, + DefaultParams: func() map[string]any { return map[string]any{} }, + }, + } +} diff --git a/internal/service/reqlog/logger.go b/internal/service/reqlog/logger.go new file mode 100644 index 0000000..414b7b5 --- /dev/null +++ b/internal/service/reqlog/logger.go @@ -0,0 +1,84 @@ +package reqlog + +import ( + "context" + "strings" + "time" + + "github.com/rose_cat707/Atlas/internal/store" +) + +type Logger struct { + store *store.Store +} + +func New(st *store.Store) *Logger { + return &Logger{store: st} +} + +type Entry struct { + SessionID string + Method string + Path string + Protocol string + ProviderID *int64 + Model string + StatusCode int + LatencyMS int + Error string + RequestSummary string + ResponseSummary string +} + +func (l *Logger) Log(ctx context.Context, e Entry) { + if l == nil || l.store == nil { + return + } + entry := e + l.store.WriteQueue.Enqueue(func(ctx context.Context, st *store.Store) error { + _, err := st.InsertRequestLog(ctx, &store.RequestLog{ + SessionID: entry.SessionID, + Method: entry.Method, + Path: entry.Path, + Protocol: entry.Protocol, + ProviderID: entry.ProviderID, + Model: entry.Model, + StatusCode: entry.StatusCode, + LatencyMS: entry.LatencyMS, + Error: entry.Error, + RequestSummary: entry.RequestSummary, + ResponseSummary: entry.ResponseSummary, + }) + return err + }) +} + +func Summarize(body []byte, max int) string { + s := strings.TrimSpace(string(body)) + if len(s) <= max { + return s + } + return s[:max] + "..." +} + +func SummarizeSSE(raw string, max int) string { + var b strings.Builder + for _, line := range strings.Split(raw, "\n") { + if strings.HasPrefix(line, "data: ") && !strings.Contains(line, "[DONE]") { + b.WriteString(strings.TrimPrefix(line, "data: ")) + b.WriteByte(' ') + } + } + return Summarize([]byte(strings.TrimSpace(b.String())), max) +} + +func Since(start time.Time) int { + return int(time.Since(start).Milliseconds()) +} + +func ProviderIDPtr(id int64) *int64 { + if id <= 0 { + return nil + } + return &id +} diff --git a/internal/service/taskevents/hub.go b/internal/service/taskevents/hub.go new file mode 100644 index 0000000..63420bc --- /dev/null +++ b/internal/service/taskevents/hub.go @@ -0,0 +1,58 @@ +package taskevents + +import ( + "sync" + + "github.com/rose_cat707/Atlas/internal/store" +) + +type Hub struct { + mu sync.RWMutex + subs map[int64]map[chan store.Task]struct{} +} + +func NewHub() *Hub { + return &Hub{subs: make(map[int64]map[chan store.Task]struct{})} +} + +func (h *Hub) Subscribe(taskID int64) chan store.Task { + ch := make(chan store.Task, 4) + h.mu.Lock() + if h.subs[taskID] == nil { + h.subs[taskID] = make(map[chan store.Task]struct{}) + } + h.subs[taskID][ch] = struct{}{} + h.mu.Unlock() + return ch +} + +func (h *Hub) Unsubscribe(taskID int64, ch chan store.Task) { + h.mu.Lock() + if m, ok := h.subs[taskID]; ok { + delete(m, ch) + if len(m) == 0 { + delete(h.subs, taskID) + } + } + h.mu.Unlock() +} + +func (h *Hub) Publish(task store.Task) { + h.mu.RLock() + defer h.mu.RUnlock() + for ch := range h.subs[task.ID] { + select { + case ch <- task: + default: + } + } +} + +func Terminal(status string) bool { + switch status { + case "succeeded", "failed", "canceled": + return true + default: + return false + } +} diff --git a/internal/settings/keys.go b/internal/settings/keys.go index 969d961..3d696e7 100644 --- a/internal/settings/keys.go +++ b/internal/settings/keys.go @@ -6,14 +6,16 @@ const ( KeyAdminPassword = "admin_password" KeyAuthEnabled = "auth_enabled" KeyUILocale = "ui_locale" + KeyFileSignSecret = "file_sign_secret" ) var Defaults = map[string]string{ - KeyAppName: "App", - KeyAPIPort: "8080", - KeyAdminPassword: "admin", - KeyAuthEnabled: "true", - KeyUILocale: "zh-CN", + KeyAppName: "Atlas", + KeyAPIPort: "8080", + KeyAdminPassword: "admin", + KeyAuthEnabled: "true", + KeyUILocale: "zh-CN", + KeyFileSignSecret: "change-me-atlas-file-sign", } // Display returns settings safe for GET responses (masks secrets). diff --git a/internal/store/cache.go b/internal/store/cache.go new file mode 100644 index 0000000..336cb4d --- /dev/null +++ b/internal/store/cache.go @@ -0,0 +1,46 @@ +package store + +import ( + "sync" + "time" +) + +type HotCache struct { + mu sync.RWMutex + entries map[string]cacheEntry + ttl time.Duration +} + +type cacheEntry struct { + value any + expiresAt time.Time +} + +func NewHotCache(ttl time.Duration) *HotCache { + return &HotCache{ + entries: make(map[string]cacheEntry), + ttl: ttl, + } +} + +func (c *HotCache) Get(key string) (any, bool) { + c.mu.RLock() + defer c.mu.RUnlock() + e, ok := c.entries[key] + if !ok || time.Now().After(e.expiresAt) { + return nil, false + } + return e.value, true +} + +func (c *HotCache) Set(key string, value any) { + c.mu.Lock() + defer c.mu.Unlock() + c.entries[key] = cacheEntry{value: value, expiresAt: time.Now().Add(c.ttl)} +} + +func (c *HotCache) Delete(key string) { + c.mu.Lock() + defer c.mu.Unlock() + delete(c.entries, key) +} diff --git a/internal/store/conversations.go b/internal/store/conversations.go new file mode 100644 index 0000000..f0ec12a --- /dev/null +++ b/internal/store/conversations.go @@ -0,0 +1,140 @@ +package store + +import "context" + +type Conversation struct { + ID int64 `json:"id"` + SessionID string `json:"session_id"` + Title string `json:"title"` + Model string `json:"model"` + ParamsJSON string `json:"params_json"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +type ChatMessageRow struct { + ID int64 `json:"id"` + ConversationID int64 `json:"conversation_id"` + Role string `json:"role"` + Content string `json:"content"` + CreatedAt string `json:"created_at"` +} + +func (s *Store) ListConversations(ctx context.Context, sessionID string) ([]Conversation, error) { + _ = s.PurgeEmptyConversations(ctx, sessionID) + rows, err := s.db.QueryContext(ctx, + `SELECT c.id, c.session_id, c.title, c.model, c.params_json, c.created_at, c.updated_at + FROM chat_conversations c + WHERE c.session_id = ? + AND EXISTS (SELECT 1 FROM chat_messages m WHERE m.conversation_id = c.id) + ORDER BY c.updated_at DESC`, sessionID, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Conversation + for rows.Next() { + var c Conversation + if err := rows.Scan(&c.ID, &c.SessionID, &c.Title, &c.Model, &c.ParamsJSON, &c.CreatedAt, &c.UpdatedAt); err != nil { + return nil, err + } + items = append(items, c) + } + return items, rows.Err() +} + +func (s *Store) GetConversation(ctx context.Context, sessionID string, id int64) (*Conversation, error) { + var c Conversation + err := s.db.QueryRowContext(ctx, + `SELECT id, session_id, title, model, params_json, created_at, updated_at FROM chat_conversations WHERE id = ? AND session_id = ?`, id, sessionID, + ).Scan(&c.ID, &c.SessionID, &c.Title, &c.Model, &c.ParamsJSON, &c.CreatedAt, &c.UpdatedAt) + if err != nil { + return nil, err + } + return &c, nil +} + +func (s *Store) CreateConversation(ctx context.Context, c *Conversation) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO chat_conversations (session_id, title, model, params_json) VALUES (?, ?, ?, ?)`, + c.SessionID, c.Title, c.Model, c.ParamsJSON, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) UpdateConversation(ctx context.Context, c *Conversation) error { + _, err := s.db.ExecContext(ctx, + `UPDATE chat_conversations SET title=?, model=?, params_json=?, updated_at=datetime('now') WHERE id=? AND session_id=?`, + c.Title, c.Model, c.ParamsJSON, c.ID, c.SessionID, + ) + return err +} + +func (s *Store) DeleteConversation(ctx context.Context, sessionID string, id int64) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM chat_conversations WHERE id = ? AND session_id = ?`, id, sessionID) + return err +} + +func (s *Store) PurgeEmptyConversations(ctx context.Context, sessionID string) error { + _, err := s.db.ExecContext(ctx, + `DELETE FROM chat_conversations + WHERE session_id = ? + AND NOT EXISTS (SELECT 1 FROM chat_messages m WHERE m.conversation_id = chat_conversations.id)`, + sessionID, + ) + return err +} + +func (s *Store) ListMessages(ctx context.Context, conversationID int64, limit int) ([]ChatMessageRow, error) { + if limit <= 0 { + limit = 80 + } + rows, err := s.db.QueryContext(ctx, + `SELECT id, conversation_id, role, content, created_at FROM chat_messages WHERE conversation_id = ? ORDER BY id DESC LIMIT ?`, conversationID, limit, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ChatMessageRow + for rows.Next() { + var m ChatMessageRow + if err := rows.Scan(&m.ID, &m.ConversationID, &m.Role, &m.Content, &m.CreatedAt); err != nil { + return nil, err + } + items = append(items, m) + } + for i, j := 0, len(items)-1; i < j; i, j = i+1, j-1 { + items[i], items[j] = items[j], items[i] + } + return items, rows.Err() +} + +func (s *Store) InsertMessage(ctx context.Context, m *ChatMessageRow) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO chat_messages (conversation_id, role, content) VALUES (?, ?, ?)`, + m.ConversationID, m.Role, m.Content, + ) + if err != nil { + return 0, err + } + _, _ = s.db.ExecContext(ctx, `UPDATE chat_conversations SET updated_at=datetime('now') WHERE id=?`, m.ConversationID) + return res.LastInsertId() +} + +func (s *Store) CountUserMessages(ctx context.Context, conversationID int64) (int, error) { + var count int + err := s.db.QueryRowContext(ctx, + `SELECT COUNT(*) FROM chat_messages WHERE conversation_id = ? AND role = 'user'`, conversationID, + ).Scan(&count) + return count, err +} + +func (s *Store) ClearMessages(ctx context.Context, conversationID int64) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM chat_messages WHERE conversation_id = ?`, conversationID) + return err +} diff --git a/internal/store/files.go b/internal/store/files.go new file mode 100644 index 0000000..1d05627 --- /dev/null +++ b/internal/store/files.go @@ -0,0 +1,44 @@ +package store + +import "context" + +type FileRecord struct { + ID int64 `json:"id"` + SessionID string `json:"session_id"` + Kind string `json:"kind"` + Mime string `json:"mime"` + StoragePath string `json:"storage_path"` + SizeBytes int64 `json:"size_bytes"` + SHA256 string `json:"sha256"` + ExpiresAt *string `json:"expires_at"` + CreatedAt string `json:"created_at"` +} + +func (s *Store) InsertFile(ctx context.Context, f *FileRecord) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO files (session_id, kind, mime, storage_path, size_bytes, sha256, expires_at) VALUES (?, ?, ?, ?, ?, ?, ?)`, + f.SessionID, f.Kind, f.Mime, f.StoragePath, f.SizeBytes, f.SHA256, f.ExpiresAt, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) GetFile(ctx context.Context, sessionID string, id int64) (*FileRecord, error) { + var f FileRecord + err := s.db.QueryRowContext(ctx, + `SELECT id, session_id, kind, mime, storage_path, size_bytes, sha256, expires_at, created_at FROM files WHERE id = ? AND session_id = ?`, id, sessionID, + ).Scan(&f.ID, &f.SessionID, &f.Kind, &f.Mime, &f.StoragePath, &f.SizeBytes, &f.SHA256, &f.ExpiresAt, &f.CreatedAt) + if err != nil { + return nil, err + } + return &f, nil +} + +func (s *Store) CountFiles(ctx context.Context) (int, int64, error) { + var count int + var total int64 + err := s.db.QueryRowContext(ctx, `SELECT COUNT(*), COALESCE(SUM(size_bytes),0) FROM files`).Scan(&count, &total) + return count, total, err +} diff --git a/internal/store/generations.go b/internal/store/generations.go new file mode 100644 index 0000000..78c75e1 --- /dev/null +++ b/internal/store/generations.go @@ -0,0 +1,128 @@ +package store + +import "context" + +type Generation struct { + ID int64 `json:"id"` + SessionID string `json:"session_id"` + ServiceType string `json:"service_type"` + ProviderID *int64 `json:"provider_id"` + Model string `json:"model"` + InputJSON string `json:"input_json"` + OutputJSON string `json:"output_json"` + TaskID *int64 `json:"task_id"` + CreatedAt string `json:"created_at"` +} + +func (s *Store) ListGenerations(ctx context.Context, sessionID, serviceType string, limit, offset int) ([]Generation, int, error) { + countQ := `SELECT COUNT(*) FROM generations WHERE session_id = ?` + args := []any{sessionID} + if serviceType != "" { + countQ += ` AND service_type = ?` + args = append(args, serviceType) + } + var total int + if err := s.db.QueryRowContext(ctx, countQ, args...).Scan(&total); err != nil { + return nil, 0, err + } + q := `SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE session_id = ?` + qArgs := []any{sessionID} + if serviceType != "" { + q += ` AND service_type = ?` + qArgs = append(qArgs, serviceType) + } + q += ` ORDER BY id DESC LIMIT ? OFFSET ?` + qArgs = append(qArgs, limit, offset) + rows, err := s.db.QueryContext(ctx, q, qArgs...) + if err != nil { + return nil, 0, err + } + defer rows.Close() + var items []Generation + for rows.Next() { + var g Generation + if err := rows.Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt); err != nil { + return nil, 0, err + } + items = append(items, g) + } + return items, total, rows.Err() +} + +func (s *Store) GetGeneration(ctx context.Context, sessionID string, id int64) (*Generation, error) { + var g Generation + err := s.db.QueryRowContext(ctx, + `SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE id = ? AND session_id = ?`, id, sessionID, + ).Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt) + if err != nil { + return nil, err + } + return &g, nil +} + +func (s *Store) InsertGeneration(ctx context.Context, g *Generation) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO generations (session_id, service_type, provider_id, model, input_json, output_json, task_id) VALUES (?, ?, ?, ?, ?, ?, ?)`, + g.SessionID, g.ServiceType, g.ProviderID, g.Model, g.InputJSON, g.OutputJSON, g.TaskID, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) DeleteGeneration(ctx context.Context, sessionID string, id int64) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM generations WHERE id = ? AND session_id = ?`, id, sessionID) + return err +} + +func (s *Store) ListAllGenerations(ctx context.Context, serviceType string, limit, offset int) ([]Generation, int, error) { + countQ := `SELECT COUNT(*) FROM generations WHERE 1=1` + args := []any{} + if serviceType != "" { + countQ += ` AND service_type = ?` + args = append(args, serviceType) + } + var total int + if err := s.db.QueryRowContext(ctx, countQ, args...).Scan(&total); err != nil { + return nil, 0, err + } + q := `SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE 1=1` + qArgs := []any{} + if serviceType != "" { + q += ` AND service_type = ?` + qArgs = append(qArgs, serviceType) + } + q += ` ORDER BY id DESC LIMIT ? OFFSET ?` + qArgs = append(qArgs, limit, offset) + rows, err := s.db.QueryContext(ctx, q, qArgs...) + if err != nil { + return nil, 0, err + } + defer rows.Close() + var items []Generation + for rows.Next() { + var g Generation + if err := rows.Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt); err != nil { + return nil, 0, err + } + items = append(items, g) + } + return items, total, rows.Err() +} + +func (s *Store) AdminGetGeneration(ctx context.Context, id int64) (*Generation, error) { + var g Generation + err := s.db.QueryRowContext(ctx, + `SELECT id, session_id, service_type, provider_id, model, input_json, output_json, task_id, created_at FROM generations WHERE id = ?`, id, + ).Scan(&g.ID, &g.SessionID, &g.ServiceType, &g.ProviderID, &g.Model, &g.InputJSON, &g.OutputJSON, &g.TaskID, &g.CreatedAt) + if err != nil { + return nil, err + } + return &g, nil +} + +func (s *Store) AdminDeleteGeneration(ctx context.Context, id int64) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM generations WHERE id = ?`, id) + return err +} diff --git a/internal/store/migrations/002_atlas.sql b/internal/store/migrations/002_atlas.sql new file mode 100644 index 0000000..95b53ac --- /dev/null +++ b/internal/store/migrations/002_atlas.sql @@ -0,0 +1,131 @@ +-- +goose Up +CREATE TABLE IF NOT EXISTS providers ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + name TEXT NOT NULL, + protocol TEXT NOT NULL, + base_url TEXT NOT NULL, + api_key TEXT NOT NULL DEFAULT '', + service_type TEXT NOT NULL, + is_default INTEGER NOT NULL DEFAULT 0, + enabled INTEGER NOT NULL DEFAULT 1, + extra_json TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) +); + +CREATE TABLE IF NOT EXISTS user_sessions ( + id TEXT PRIMARY KEY, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + last_seen_at TEXT NOT NULL DEFAULT (datetime('now')) +); + +CREATE TABLE IF NOT EXISTS chat_conversations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + title TEXT NOT NULL DEFAULT '新对话', + model TEXT NOT NULL DEFAULT '', + params_json TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')) +); +CREATE INDEX IF NOT EXISTS idx_chat_conversations_session ON chat_conversations(session_id); + +CREATE TABLE IF NOT EXISTS chat_messages ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + conversation_id INTEGER NOT NULL, + role TEXT NOT NULL, + content TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + FOREIGN KEY (conversation_id) REFERENCES chat_conversations(id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_chat_messages_conv ON chat_messages(conversation_id); + +CREATE TABLE IF NOT EXISTS generations ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + service_type TEXT NOT NULL, + provider_id INTEGER, + model TEXT NOT NULL DEFAULT '', + input_json TEXT NOT NULL DEFAULT '{}', + output_json TEXT NOT NULL DEFAULT '{}', + task_id INTEGER, + created_at TEXT NOT NULL DEFAULT (datetime('now')) +); +CREATE INDEX IF NOT EXISTS idx_generations_session ON generations(session_id, created_at DESC); + +CREATE TABLE IF NOT EXISTS presets ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + name TEXT NOT NULL, + service_type TEXT NOT NULL, + params_json TEXT NOT NULL DEFAULT '{}', + created_at TEXT NOT NULL DEFAULT (datetime('now')) +); +CREATE INDEX IF NOT EXISTS idx_presets_session ON presets(session_id, service_type); + +CREATE TABLE IF NOT EXISTS tasks ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + service_type TEXT NOT NULL, + protocol TEXT NOT NULL, + provider_id INTEGER NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', + progress INTEGER NOT NULL DEFAULT 0, + input_json TEXT NOT NULL DEFAULT '{}', + result_json TEXT NOT NULL DEFAULT '{}', + error TEXT NOT NULL DEFAULT '', + external_id TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL DEFAULT (datetime('now')), + updated_at TEXT NOT NULL DEFAULT (datetime('now')), + completed_at TEXT +); +CREATE INDEX IF NOT EXISTS idx_tasks_session_status ON tasks(session_id, status); + +CREATE TABLE IF NOT EXISTS files ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + kind TEXT NOT NULL, + mime TEXT NOT NULL DEFAULT '', + storage_path TEXT NOT NULL, + size_bytes INTEGER NOT NULL DEFAULT 0, + sha256 TEXT NOT NULL DEFAULT '', + expires_at TEXT, + created_at TEXT NOT NULL DEFAULT (datetime('now')) +); +CREATE INDEX IF NOT EXISTS idx_files_session ON files(session_id); + +CREATE TABLE IF NOT EXISTS request_logs ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL DEFAULT '', + method TEXT NOT NULL, + path TEXT NOT NULL, + protocol TEXT NOT NULL DEFAULT '', + provider_id INTEGER, + model TEXT NOT NULL DEFAULT '', + status_code INTEGER NOT NULL DEFAULT 0, + latency_ms INTEGER NOT NULL DEFAULT 0, + error TEXT NOT NULL DEFAULT '', + request_summary TEXT NOT NULL DEFAULT '', + response_summary TEXT NOT NULL DEFAULT '', + created_at TEXT NOT NULL DEFAULT (datetime('now')) +); +CREATE INDEX IF NOT EXISTS idx_request_logs_created ON request_logs(created_at DESC); + +INSERT OR IGNORE INTO settings (key, value) VALUES + ('file_sign_secret', 'change-me-atlas-file-sign'), + ('write_flush_interval_ms', '200'), + ('task_progress_flush_interval_ms', '1000'), + ('cache_ttl_seconds', '30'); + +UPDATE settings SET value = 'Atlas' WHERE key = 'app_name'; + +-- +goose Down +DROP TABLE IF EXISTS request_logs; +DROP TABLE IF EXISTS files; +DROP TABLE IF EXISTS tasks; +DROP TABLE IF EXISTS presets; +DROP TABLE IF EXISTS generations; +DROP TABLE IF EXISTS chat_messages; +DROP TABLE IF EXISTS chat_conversations; +DROP TABLE IF EXISTS user_sessions; +DROP TABLE IF EXISTS providers; diff --git a/internal/store/migrations/003_provider_models.sql b/internal/store/migrations/003_provider_models.sql new file mode 100644 index 0000000..f1bb496 --- /dev/null +++ b/internal/store/migrations/003_provider_models.sql @@ -0,0 +1,14 @@ +-- +goose Up +CREATE TABLE IF NOT EXISTS provider_models ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + provider_id INTEGER NOT NULL, + model_id TEXT NOT NULL, + name TEXT NOT NULL, + created_at TEXT NOT NULL DEFAULT (datetime('now')), + UNIQUE(provider_id, model_id), + FOREIGN KEY (provider_id) REFERENCES providers(id) ON DELETE CASCADE +); +CREATE INDEX IF NOT EXISTS idx_provider_models_provider ON provider_models(provider_id); + +-- +goose Down +DROP TABLE IF EXISTS provider_models; diff --git a/internal/store/presets.go b/internal/store/presets.go new file mode 100644 index 0000000..7e7d8cf --- /dev/null +++ b/internal/store/presets.go @@ -0,0 +1,60 @@ +package store + +import "context" + +type Preset struct { + ID int64 `json:"id"` + SessionID string `json:"session_id"` + Name string `json:"name"` + ServiceType string `json:"service_type"` + ParamsJSON string `json:"params_json"` + CreatedAt string `json:"created_at"` +} + +func (s *Store) ListPresets(ctx context.Context, sessionID, serviceType string) ([]Preset, error) { + q := `SELECT id, session_id, name, service_type, params_json, created_at FROM presets WHERE session_id = ?` + args := []any{sessionID} + if serviceType != "" { + q += ` AND service_type = ?` + args = append(args, serviceType) + } + q += ` ORDER BY id DESC` + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Preset + for rows.Next() { + var p Preset + if err := rows.Scan(&p.ID, &p.SessionID, &p.Name, &p.ServiceType, &p.ParamsJSON, &p.CreatedAt); err != nil { + return nil, err + } + items = append(items, p) + } + return items, rows.Err() +} + +func (s *Store) CreatePreset(ctx context.Context, p *Preset) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO presets (session_id, name, service_type, params_json) VALUES (?, ?, ?, ?)`, + p.SessionID, p.Name, p.ServiceType, p.ParamsJSON, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) UpdatePreset(ctx context.Context, p *Preset) error { + _, err := s.db.ExecContext(ctx, + `UPDATE presets SET name=?, params_json=? WHERE id=? AND session_id=?`, + p.Name, p.ParamsJSON, p.ID, p.SessionID, + ) + return err +} + +func (s *Store) DeletePreset(ctx context.Context, sessionID string, id int64) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM presets WHERE id = ? AND session_id = ?`, id, sessionID) + return err +} diff --git a/internal/store/provider_models.go b/internal/store/provider_models.go new file mode 100644 index 0000000..b67b494 --- /dev/null +++ b/internal/store/provider_models.go @@ -0,0 +1,53 @@ +package store + +import "context" + +type ProviderModel struct { + ID int64 `json:"id"` + ProviderID int64 `json:"provider_id"` + ModelID string `json:"model_id"` + Name string `json:"name"` + CreatedAt string `json:"created_at"` +} + +func (s *Store) ListProviderModels(ctx context.Context, providerID int64) ([]ProviderModel, error) { + rows, err := s.db.QueryContext(ctx, + `SELECT id, provider_id, model_id, name, created_at FROM provider_models WHERE provider_id = ? ORDER BY id ASC`, providerID, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []ProviderModel + for rows.Next() { + var m ProviderModel + if err := rows.Scan(&m.ID, &m.ProviderID, &m.ModelID, &m.Name, &m.CreatedAt); err != nil { + return nil, err + } + items = append(items, m) + } + return items, rows.Err() +} + +func (s *Store) CreateProviderModel(ctx context.Context, m *ProviderModel) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO provider_models (provider_id, model_id, name) VALUES (?, ?, ?)`, + m.ProviderID, m.ModelID, m.Name, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) DeleteProviderModel(ctx context.Context, providerID int64, modelID string) error { + _, err := s.db.ExecContext(ctx, + `DELETE FROM provider_models WHERE provider_id = ? AND model_id = ?`, providerID, modelID, + ) + return err +} + +func (s *Store) DeleteProviderModelsByProvider(ctx context.Context, providerID int64) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM provider_models WHERE provider_id = ?`, providerID) + return err +} diff --git a/internal/store/providers.go b/internal/store/providers.go new file mode 100644 index 0000000..92b83b6 --- /dev/null +++ b/internal/store/providers.go @@ -0,0 +1,136 @@ +package store + +import ( + "context" + "database/sql" + "fmt" + + "github.com/rose_cat707/Atlas/internal/domain" +) + +type Provider struct { + ID int64 `json:"id"` + Name string `json:"name"` + Protocol string `json:"protocol"` + BaseURL string `json:"base_url"` + APIKey string `json:"api_key,omitempty"` + ServiceType string `json:"service_type"` + IsDefault bool `json:"is_default"` + Enabled bool `json:"enabled"` + ExtraJSON string `json:"extra_json"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` +} + +func (s *Store) ListProviders(ctx context.Context, serviceType string, enabledOnly bool) ([]Provider, error) { + q := `SELECT id, name, protocol, base_url, api_key, service_type, is_default, enabled, extra_json, created_at, updated_at FROM providers WHERE 1=1` + args := []any{} + if serviceType != "" { + q += ` AND service_type = ?` + args = append(args, serviceType) + } + if enabledOnly { + q += ` AND enabled = 1` + } + q += ` ORDER BY is_default DESC, id ASC` + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Provider + for rows.Next() { + var p Provider + var isDef, en int + if err := rows.Scan(&p.ID, &p.Name, &p.Protocol, &p.BaseURL, &p.APIKey, &p.ServiceType, &isDef, &en, &p.ExtraJSON, &p.CreatedAt, &p.UpdatedAt); err != nil { + return nil, err + } + p.IsDefault = isDef == 1 + p.Enabled = en == 1 + items = append(items, p) + } + return items, rows.Err() +} + +func (s *Store) GetProvider(ctx context.Context, id int64) (*Provider, error) { + var p Provider + var isDef, en int + err := s.db.QueryRowContext(ctx, + `SELECT id, name, protocol, base_url, api_key, service_type, is_default, enabled, extra_json, created_at, updated_at FROM providers WHERE id = ?`, id, + ).Scan(&p.ID, &p.Name, &p.Protocol, &p.BaseURL, &p.APIKey, &p.ServiceType, &isDef, &en, &p.ExtraJSON, &p.CreatedAt, &p.UpdatedAt) + if err != nil { + return nil, err + } + p.IsDefault = isDef == 1 + p.Enabled = en == 1 + return &p, nil +} + +func (s *Store) CreateProvider(ctx context.Context, p *Provider) (int64, error) { + if p.IsDefault { + if _, err := s.db.ExecContext(ctx, `UPDATE providers SET is_default = 0 WHERE service_type = ?`, p.ServiceType); err != nil { + return 0, err + } + } + res, err := s.db.ExecContext(ctx, + `INSERT INTO providers (name, protocol, base_url, api_key, service_type, is_default, enabled, extra_json) VALUES (?, ?, ?, ?, ?, ?, ?, ?)`, + p.Name, p.Protocol, p.BaseURL, p.APIKey, p.ServiceType, boolInt(p.IsDefault), boolInt(p.Enabled), p.ExtraJSON, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) UpdateProvider(ctx context.Context, p *Provider) error { + if p.IsDefault { + if _, err := s.db.ExecContext(ctx, `UPDATE providers SET is_default = 0 WHERE service_type = ? AND id != ?`, p.ServiceType, p.ID); err != nil { + return err + } + } + _, err := s.db.ExecContext(ctx, + `UPDATE providers SET name=?, protocol=?, base_url=?, api_key=?, service_type=?, is_default=?, enabled=?, extra_json=?, updated_at=datetime('now') WHERE id=?`, + p.Name, p.Protocol, p.BaseURL, p.APIKey, p.ServiceType, boolInt(p.IsDefault), boolInt(p.Enabled), p.ExtraJSON, p.ID, + ) + return err +} + +func (s *Store) DeleteProvider(ctx context.Context, id int64) error { + _, err := s.db.ExecContext(ctx, `DELETE FROM providers WHERE id = ?`, id) + return err +} + +func boolInt(v bool) int { + if v { + return 1 + } + return 0 +} + +func (s *Store) DefaultProvider(ctx context.Context, serviceType string) (*Provider, error) { + providers, err := s.ListProviders(ctx, serviceType, true) + if err != nil { + return nil, err + } + for _, p := range providers { + if p.IsDefault { + return &p, nil + } + } + if len(providers) > 0 { + return &providers[0], nil + } + return nil, sql.ErrNoRows +} + +func MaskProvider(p Provider) Provider { + p.APIKey = MaskSecret(p.APIKey) + return p +} + +func ValidateProvider(p *Provider) error { + if p.Name == "" || p.Protocol == "" || p.BaseURL == "" || p.ServiceType == "" { + return fmt.Errorf("name, protocol, base_url, service_type required") + } + return domain.ValidateProviderBinding(p.Protocol, p.ServiceType) +} diff --git a/internal/store/request_logs.go b/internal/store/request_logs.go new file mode 100644 index 0000000..38d321b --- /dev/null +++ b/internal/store/request_logs.go @@ -0,0 +1,101 @@ +package store + +import "context" + +type RequestLog struct { + ID int64 `json:"id"` + SessionID string `json:"session_id"` + Method string `json:"method"` + Path string `json:"path"` + Protocol string `json:"protocol"` + ProviderID *int64 `json:"provider_id"` + Model string `json:"model"` + StatusCode int `json:"status_code"` + LatencyMS int `json:"latency_ms"` + Error string `json:"error"` + RequestSummary string `json:"request_summary"` + ResponseSummary string `json:"response_summary"` + CreatedAt string `json:"created_at"` +} + +func (s *Store) InsertRequestLog(ctx context.Context, l *RequestLog) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO request_logs (session_id, method, path, protocol, provider_id, model, status_code, latency_ms, error, request_summary, response_summary) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + l.SessionID, l.Method, l.Path, l.Protocol, l.ProviderID, l.Model, l.StatusCode, l.LatencyMS, l.Error, l.RequestSummary, l.ResponseSummary, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) ListRequestLogs(ctx context.Context, filter RequestLogFilter, limit, offset int) ([]RequestLog, int, error) { + where := `WHERE 1=1` + args := []any{} + if filter.Path != "" { + where += ` AND path LIKE ?` + args = append(args, "%"+filter.Path+"%") + } + if filter.StatusCode > 0 { + where += ` AND status_code = ?` + args = append(args, filter.StatusCode) + } + if filter.ProviderID > 0 { + where += ` AND provider_id = ?` + args = append(args, filter.ProviderID) + } + var total int + if err := s.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM request_logs `+where, args...).Scan(&total); err != nil { + return nil, 0, err + } + q := `SELECT id, session_id, method, path, protocol, provider_id, model, status_code, latency_ms, error, request_summary, response_summary, created_at FROM request_logs ` + where + ` ORDER BY id DESC LIMIT ? OFFSET ?` + qArgs := append(args, limit, offset) + rows, err := s.db.QueryContext(ctx, q, qArgs...) + if err != nil { + return nil, 0, err + } + defer rows.Close() + var items []RequestLog + for rows.Next() { + var l RequestLog + if err := rows.Scan(&l.ID, &l.SessionID, &l.Method, &l.Path, &l.Protocol, &l.ProviderID, &l.Model, &l.StatusCode, &l.LatencyMS, &l.Error, &l.RequestSummary, &l.ResponseSummary, &l.CreatedAt); err != nil { + return nil, 0, err + } + items = append(items, l) + } + return items, total, rows.Err() +} + +type RequestLogFilter struct { + Path string + StatusCode int + ProviderID int64 +} + +type DashboardStats struct { + TotalRequests int `json:"total_requests"` + SuccessRate float64 `json:"success_rate"` + AvgLatencyMS float64 `json:"avg_latency_ms"` + P50LatencyMS float64 `json:"p50_latency_ms"` + P95LatencyMS float64 `json:"p95_latency_ms"` + ByServiceType map[string]int `json:"by_service_type"` + ByProvider map[string]int `json:"by_provider"` + FileCount int `json:"file_count"` + FileBytes int64 `json:"file_bytes"` +} + +func (s *Store) DashboardStats(ctx context.Context) (*DashboardStats, error) { + stats := &DashboardStats{ByServiceType: map[string]int{}, ByProvider: map[string]int{}} + var success int + err := s.db.QueryRowContext(ctx, + `SELECT COUNT(*), COALESCE(AVG(latency_ms),0), SUM(CASE WHEN status_code >= 200 AND status_code < 400 THEN 1 ELSE 0 END) FROM request_logs WHERE created_at >= datetime('now','-7 days')`, + ).Scan(&stats.TotalRequests, &stats.AvgLatencyMS, &success) + if err != nil { + return nil, err + } + if stats.TotalRequests > 0 { + stats.SuccessRate = float64(success) / float64(stats.TotalRequests) + } + stats.FileCount, stats.FileBytes, _ = s.CountFiles(ctx) + return stats, nil +} diff --git a/internal/store/sessions.go b/internal/store/sessions.go new file mode 100644 index 0000000..bd9f352 --- /dev/null +++ b/internal/store/sessions.go @@ -0,0 +1,21 @@ +package store + +import "context" + +type UserSession struct { + ID string `json:"id"` + CreatedAt string `json:"created_at"` + LastSeenAt string `json:"last_seen_at"` +} + +func (s *Store) UpsertSession(ctx context.Context, id string) error { + _, err := s.db.ExecContext(ctx, + `INSERT INTO user_sessions (id, last_seen_at) VALUES (?, datetime('now')) + ON CONFLICT(id) DO UPDATE SET last_seen_at = datetime('now')`, id, + ) + return err +} + +func (s *Store) TouchSession(ctx context.Context, id string) error { + return s.UpsertSession(ctx, id) +} diff --git a/internal/store/store.go b/internal/store/store.go index 4fee5d0..a06364c 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -7,6 +7,7 @@ import ( "fmt" "os" "path/filepath" + "time" "github.com/pressly/goose/v3" _ "modernc.org/sqlite" @@ -16,7 +17,9 @@ import ( var migrationsFS embed.FS type Store struct { - db *sql.DB + db *sql.DB + Cache *HotCache + WriteQueue *WriteQueue } func Open(dataDir string) (*Store, error) { @@ -34,12 +37,19 @@ func Open(dataDir string) (*Store, error) { db.Close() return nil, err } + s.Cache = NewHotCache(30 * time.Second) + s.WriteQueue = NewWriteQueue(s, 200*time.Millisecond) return s, nil } func (s *Store) DB() *sql.DB { return s.db } -func (s *Store) Close() error { return s.db.Close() } +func (s *Store) Close() error { + if s.WriteQueue != nil { + s.WriteQueue.Close() + } + return s.db.Close() +} func (s *Store) migrate() error { goose.SetBaseFS(migrationsFS) @@ -143,7 +153,8 @@ func MaskSecret(value string) string { func ValidateSettingKey(key string) error { switch key { - case "app_name", "api_port", "admin_password", "auth_enabled", "ui_locale": + case "app_name", "api_port", "admin_password", "auth_enabled", "ui_locale", + "file_sign_secret", "write_flush_interval_ms", "task_progress_flush_interval_ms", "cache_ttl_seconds": return nil default: return fmt.Errorf("unknown setting: %s", key) diff --git a/internal/store/tasks.go b/internal/store/tasks.go new file mode 100644 index 0000000..b334f14 --- /dev/null +++ b/internal/store/tasks.go @@ -0,0 +1,93 @@ +package store + +import "context" + +type Task struct { + ID int64 `json:"id"` + SessionID string `json:"session_id"` + ServiceType string `json:"service_type"` + Protocol string `json:"protocol"` + ProviderID int64 `json:"provider_id"` + Status string `json:"status"` + Progress int `json:"progress"` + InputJSON string `json:"input_json"` + ResultJSON string `json:"result_json"` + Error string `json:"error"` + ExternalID string `json:"external_id"` + CreatedAt string `json:"created_at"` + UpdatedAt string `json:"updated_at"` + CompletedAt *string `json:"completed_at"` +} + +func (s *Store) InsertTask(ctx context.Context, t *Task) (int64, error) { + res, err := s.db.ExecContext(ctx, + `INSERT INTO tasks (session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + t.SessionID, t.ServiceType, t.Protocol, t.ProviderID, t.Status, t.Progress, t.InputJSON, t.ResultJSON, t.Error, t.ExternalID, + ) + if err != nil { + return 0, err + } + return res.LastInsertId() +} + +func (s *Store) UpdateTask(ctx context.Context, t *Task) error { + _, err := s.db.ExecContext(ctx, + `UPDATE tasks SET status=?, progress=?, result_json=?, error=?, external_id=?, updated_at=datetime('now'), completed_at=? WHERE id=? AND session_id=?`, + t.Status, t.Progress, t.ResultJSON, t.Error, t.ExternalID, t.CompletedAt, t.ID, t.SessionID, + ) + return err +} + +func (s *Store) GetTask(ctx context.Context, sessionID string, id int64) (*Task, error) { + var t Task + err := s.db.QueryRowContext(ctx, + `SELECT id, session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id, created_at, updated_at, completed_at FROM tasks WHERE id = ? AND session_id = ?`, id, sessionID, + ).Scan(&t.ID, &t.SessionID, &t.ServiceType, &t.Protocol, &t.ProviderID, &t.Status, &t.Progress, &t.InputJSON, &t.ResultJSON, &t.Error, &t.ExternalID, &t.CreatedAt, &t.UpdatedAt, &t.CompletedAt) + if err != nil { + return nil, err + } + return &t, nil +} + +func (s *Store) ListActiveTasks(ctx context.Context, sessionID, serviceType string) ([]Task, error) { + q := `SELECT id, session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id, created_at, updated_at, completed_at FROM tasks WHERE session_id = ? AND status IN ('pending','running')` + args := []any{sessionID} + if serviceType != "" { + q += ` AND service_type = ?` + args = append(args, serviceType) + } + q += ` ORDER BY id DESC` + rows, err := s.db.QueryContext(ctx, q, args...) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Task + for rows.Next() { + var t Task + if err := rows.Scan(&t.ID, &t.SessionID, &t.ServiceType, &t.Protocol, &t.ProviderID, &t.Status, &t.Progress, &t.InputJSON, &t.ResultJSON, &t.Error, &t.ExternalID, &t.CreatedAt, &t.UpdatedAt, &t.CompletedAt); err != nil { + return nil, err + } + items = append(items, t) + } + return items, rows.Err() +} + +func (s *Store) ListPendingTasks(ctx context.Context) ([]Task, error) { + rows, err := s.db.QueryContext(ctx, + `SELECT id, session_id, service_type, protocol, provider_id, status, progress, input_json, result_json, error, external_id, created_at, updated_at, completed_at FROM tasks WHERE status IN ('pending','running') ORDER BY id ASC`, + ) + if err != nil { + return nil, err + } + defer rows.Close() + var items []Task + for rows.Next() { + var t Task + if err := rows.Scan(&t.ID, &t.SessionID, &t.ServiceType, &t.Protocol, &t.ProviderID, &t.Status, &t.Progress, &t.InputJSON, &t.ResultJSON, &t.Error, &t.ExternalID, &t.CreatedAt, &t.UpdatedAt, &t.CompletedAt); err != nil { + return nil, err + } + items = append(items, t) + } + return items, rows.Err() +} diff --git a/internal/store/write_queue.go b/internal/store/write_queue.go new file mode 100644 index 0000000..805e7c5 --- /dev/null +++ b/internal/store/write_queue.go @@ -0,0 +1,90 @@ +package store + +import ( + "context" + "log" + "sync" + "time" +) + +type writeJob func(ctx context.Context, db *Store) error + +type WriteQueue struct { + store *Store + ch chan writeJob + wg sync.WaitGroup + interval time.Duration + stopCh chan struct{} +} + +func NewWriteQueue(s *Store, flushInterval time.Duration) *WriteQueue { + wq := &WriteQueue{ + store: s, + ch: make(chan writeJob, 256), + interval: flushInterval, + stopCh: make(chan struct{}), + } + wq.wg.Add(1) + go wq.run() + return wq +} + +func (wq *WriteQueue) Enqueue(job writeJob) { + select { + case wq.ch <- job: + default: + go func() { wq.ch <- job }() + } +} + +func (wq *WriteQueue) ForceFlush(ctx context.Context, job writeJob) error { + return job(ctx, wq.store) +} + +func (wq *WriteQueue) run() { + defer wq.wg.Done() + ticker := time.NewTicker(wq.interval) + defer ticker.Stop() + pending := make([]writeJob, 0, 32) + flush := func() { + if len(pending) == 0 { + return + } + jobs := pending + pending = pending[:0] + ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) + defer cancel() + for _, job := range jobs { + if err := job(ctx, wq.store); err != nil { + log.Printf("write queue: %v", err) + } + } + } + for { + select { + case <-wq.stopCh: + flush() + for { + select { + case job := <-wq.ch: + pending = append(pending, job) + default: + flush() + return + } + } + case job := <-wq.ch: + pending = append(pending, job) + if len(pending) >= 32 { + flush() + } + case <-ticker.C: + flush() + } + } +} + +func (wq *WriteQueue) Close() { + close(wq.stopCh) + wq.wg.Wait() +} diff --git a/web/index.html b/web/index.html index a45d70a..f77ecfb 100644 --- a/web/index.html +++ b/web/index.html @@ -2,8 +2,10 @@ + - App + + Atlas
diff --git a/web/public/favicon.svg b/web/public/favicon.svg new file mode 100644 index 0000000..7fd2100 --- /dev/null +++ b/web/public/favicon.svg @@ -0,0 +1,7 @@ + + + + + + + diff --git a/web/src/api/client.ts b/web/src/api/client.ts index 963b5d4..4dc5485 100644 --- a/web/src/api/client.ts +++ b/web/src/api/client.ts @@ -19,8 +19,8 @@ api.interceptors.response.use( (err) => { if (err.response?.status === 401 && err.config?.url !== '/auth/login') { localStorage.removeItem('app_token') - if (router.currentRoute.value.name !== 'login') { - router.push('/login') + if (!router.currentRoute.value.path.startsWith('/admin/login')) { + router.push('/admin/login') } } return Promise.reject(err) diff --git a/web/src/components/AppCard.vue b/web/src/components/AppCard.vue index 998ac70..b3b95d3 100644 --- a/web/src/components/AppCard.vue +++ b/web/src/components/AppCard.vue @@ -1,5 +1,16 @@ + + diff --git a/web/src/components/AtlasLogo.vue b/web/src/components/AtlasLogo.vue new file mode 100644 index 0000000..dbe045c --- /dev/null +++ b/web/src/components/AtlasLogo.vue @@ -0,0 +1,43 @@ + + + + + diff --git a/web/src/components/PageHeader.vue b/web/src/components/PageHeader.vue index 2959d46..b200513 100644 --- a/web/src/components/PageHeader.vue +++ b/web/src/components/PageHeader.vue @@ -1,14 +1,14 @@ \ No newline at end of file diff --git a/web/src/views/TextToImageView.vue b/web/src/views/TextToImageView.vue new file mode 100644 index 0000000..19e5472 --- /dev/null +++ b/web/src/views/TextToImageView.vue @@ -0,0 +1,17 @@ + + +