初始提交：边缘AI算力机统一AI通讯层

2026-08-03 07:44:05 +08:00
commit 93a469061d
51 changed files with 11565 additions and 0 deletions
@@ -0,0 +1,99 @@
+package adapter
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	"github.com/edgeai/gateway/pkg/api"
+)
+
+// ModelAdapter is the interface that all inference engine adapters must implement.
+type ModelAdapter interface {
+	// Name returns the adapter name (e.g., "ollama", "vllm").
+	Name() string
+
+	// ChatCompletion sends a non-streaming chat completion request.
+	ChatCompletion(ctx context.Context, req *ChatRequest) (*ChatResponse, error)
+
+	// ChatCompletionStream sends a streaming chat completion request.
+	ChatCompletionStream(ctx context.Context, req *ChatRequest) (<-chan StreamChunk, error)
+
+	// ListModels returns available models from the engine.
+	ListModels(ctx context.Context) ([]ModelInfo, error)
+
+	// HealthCheck checks if the engine is reachable.
+	HealthCheck(ctx context.Context) error
+
+	// Cancel cancels an in-progress request by request ID.
+	Cancel(requestID string) error
+}
+
+// ChatRequest is the internal request sent to an adapter.
+type ChatRequest struct {
+	RequestID     string
+	Model         string // actual model name
+	Messages      []api.Message
+	MaxTokens     int
+	Temperature   *float64
+	TopP          *float64
+	Stream        bool
+	CancelCh      <-chan struct{}
+}
+
+// ChatResponse is the internal response from an adapter.
+type ChatResponse struct {
+	Content       string
+	FinishReason  string
+	InputTokens   int
+	OutputTokens  int
+	ActualModel   string
+}
+
+// StreamChunk represents a single chunk in a streaming response.
+type StreamChunk struct {
+	Delta         string
+	FinishReason  string
+	InputTokens   int
+	OutputTokens  int
+	Error         error
+	Done          bool
+}
+
+// ModelInfo describes a model available in the engine.
+type ModelInfo struct {
+	Name          string
+	ContextWindow int
+}
+
+// Registry manages model adapters by provider name.
+type Registry struct {
+	adapters map[string]ModelAdapter
+}
+
+func NewRegistry() *Registry {
+	return &Registry{adapters: make(map[string]ModelAdapter)}
+}
+
+func (r *Registry) Register(name string, adapter ModelAdapter) {
+	r.adapters[name] = adapter
+}
+
+func (r *Registry) Get(name string) (ModelAdapter, error) {
+	a, ok := r.adapters[name]
+	if !ok {
+		return nil, fmt.Errorf("adapter not found: %s", name)
+	}
+	return a, nil
+}
+
+func (r *Registry) Names() []string {
+	names := make([]string, 0, len(r.adapters))
+	for n := range r.adapters {
+		names = append(names, n)
+	}
+	return names
+}
+
+// Ensure io is imported for future use (streaming readers).
+var _ = io.EOF
@@ -0,0 +1,259 @@
+package adapter
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+)
+
+// OllamaAdapter implements ModelAdapter for Ollama inference engine.
+type OllamaAdapter struct {
+	endpoint   string
+	httpClient *http.Client
+}
+
+// NewOllamaAdapter creates a new Ollama adapter.
+func NewOllamaAdapter(endpoint string) *OllamaAdapter {
+	return &OllamaAdapter{
+		endpoint: strings.TrimRight(endpoint, "/"),
+		httpClient: &http.Client{
+			Timeout: 120 * time.Second,
+		},
+	}
+}
+
+func (a *OllamaAdapter) Name() string {
+	return "ollama"
+}
+
+// ollamaChatRequest is the Ollama /api/chat request format.
+type ollamaChatRequest struct {
+	Model    string        `json:"model"`
+	Messages []ollamaMsg   `json:"messages"`
+	Stream   bool          `json:"stream"`
+	Options  ollamaOptions `json:"options,omitempty"`
+}
+
+type ollamaMsg struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type ollamaOptions struct {
+	Temperature float64 `json:"temperature,omitempty"`
+	TopP        float64 `json:"top_p,omitempty"`
+	NumPredict  int     `json:"num_predict,omitempty"`
+}
+
+// ollamaChatResponse is the Ollama /api/chat non-streaming response.
+type ollamaChatResponse struct {
+	Model     string `json:"model"`
+	Message   ollamaMsg `json:"message"`
+	Done      bool   `json:"done"`
+	PromptEvalCount int `json:"prompt_eval_count"`
+	EvalCount      int `json:"eval_count"`
+}
+
+// ollamaChatStreamResponse is a single chunk in Ollama streaming response.
+type ollamaChatStreamResponse struct {
+	Model     string `json:"model"`
+	Message   ollamaMsg `json:"message"`
+	Done      bool   `json:"done"`
+	PromptEvalCount int `json:"prompt_eval_count,omitempty"`
+	EvalCount      int `json:"eval_count,omitempty"`
+}
+
+func (a *OllamaAdapter) ChatCompletion(ctx context.Context, req *ChatRequest) (*ChatResponse, error) {
+	ollamaReq := a.buildRequest(req, false)
+
+	body, err := json.Marshal(ollamaReq)
+	if err != nil {
+		return nil, fmt.Errorf("marshal ollama request: %w", err)
+	}
+
+	httpReq, err := http.NewRequestWithContext(ctx, "POST", a.endpoint+"/api/chat", bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("create ollama request: %w", err)
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+
+	resp, err := a.httpClient.Do(httpReq)
+	if err != nil {
+		return nil, fmt.Errorf("ollama request failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		return nil, fmt.Errorf("ollama returned status %d: %s", resp.StatusCode, string(bodyBytes))
+	}
+
+	var ollamaResp ollamaChatResponse
+	if err := json.NewDecoder(resp.Body).Decode(&ollamaResp); err != nil {
+		return nil, fmt.Errorf("decode ollama response: %w", err)
+	}
+
+	return &ChatResponse{
+		Content:      ollamaResp.Message.Content,
+		FinishReason: "stop",
+		InputTokens:  ollamaResp.PromptEvalCount,
+		OutputTokens: ollamaResp.EvalCount,
+		ActualModel:  ollamaResp.Model,
+	}, nil
+}
+
+func (a *OllamaAdapter) ChatCompletionStream(ctx context.Context, req *ChatRequest) (<-chan StreamChunk, error) {
+	ollamaReq := a.buildRequest(req, true)
+
+	body, err := json.Marshal(ollamaReq)
+	if err != nil {
+		return nil, fmt.Errorf("marshal ollama stream request: %w", err)
+	}
+
+	httpReq, err := http.NewRequestWithContext(ctx, "POST", a.endpoint+"/api/chat", bytes.NewReader(body))
+	if err != nil {
+		return nil, fmt.Errorf("create ollama stream request: %w", err)
+	}
+	httpReq.Header.Set("Content-Type", "application/json")
+
+	resp, err := a.httpClient.Do(httpReq)
+	if err != nil {
+		return nil, fmt.Errorf("ollama stream request failed: %w", err)
+	}
+
+	if resp.StatusCode != http.StatusOK {
+		bodyBytes, _ := io.ReadAll(resp.Body)
+		resp.Body.Close()
+		return nil, fmt.Errorf("ollama stream returned status %d: %s", resp.StatusCode, string(bodyBytes))
+	}
+
+	ch := make(chan StreamChunk, 100)
+	go func() {
+		defer close(ch)
+		defer resp.Body.Close()
+
+		decoder := json.NewDecoder(resp.Body)
+		for {
+			var chunk ollamaChatStreamResponse
+			if err := decoder.Decode(&chunk); err != nil {
+				if err == io.EOF {
+					ch <- StreamChunk{Done: true, FinishReason: "stop"}
+					return
+				}
+				ch <- StreamChunk{Error: fmt.Errorf("decode stream chunk: %w", err)}
+				return
+			}
+
+			// Check for cancellation
+			select {
+			case <-req.CancelCh:
+				ch <- StreamChunk{Done: true, FinishReason: "cancelled"}
+				return
+			default:
+			}
+
+			if chunk.Done {
+				ch <- StreamChunk{
+					Done:         true,
+					FinishReason: "stop",
+					InputTokens:  chunk.PromptEvalCount,
+					OutputTokens: chunk.EvalCount,
+				}
+				return
+			}
+
+			if chunk.Message.Content != "" {
+				ch <- StreamChunk{Delta: chunk.Message.Content}
+			}
+		}
+	}()
+
+	return ch, nil
+}
+
+func (a *OllamaAdapter) ListModels(ctx context.Context) ([]ModelInfo, error) {
+	httpReq, err := http.NewRequestWithContext(ctx, "GET", a.endpoint+"/api/tags", nil)
+	if err != nil {
+		return nil, fmt.Errorf("create list models request: %w", err)
+	}
+
+	resp, err := a.httpClient.Do(httpReq)
+	if err != nil {
+		return nil, fmt.Errorf("list models failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("list models returned status %d", resp.StatusCode)
+	}
+
+	var tagsResp struct {
+		Models []struct {
+			Name string `json:"name"`
+		} `json:"models"`
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&tagsResp); err != nil {
+		return nil, fmt.Errorf("decode tags response: %w", err)
+	}
+
+	models := make([]ModelInfo, len(tagsResp.Models))
+	for i, m := range tagsResp.Models {
+		models[i] = ModelInfo{Name: m.Name}
+	}
+	return models, nil
+}
+
+func (a *OllamaAdapter) HealthCheck(ctx context.Context) error {
+	httpReq, err := http.NewRequestWithContext(ctx, "GET", a.endpoint+"/api/tags", nil)
+	if err != nil {
+		return fmt.Errorf("create health check request: %w", err)
+	}
+
+	resp, err := a.httpClient.Do(httpReq)
+	if err != nil {
+		return fmt.Errorf("health check failed: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("health check returned status %d", resp.StatusCode)
+	}
+	return nil
+}
+
+func (a *OllamaAdapter) Cancel(requestID string) error {
+	// Ollama doesn't support request cancellation by ID in the API.
+	// Cancellation is handled by closing the HTTP connection (context cancellation).
+	return nil
+}
+
+func (a *OllamaAdapter) buildRequest(req *ChatRequest, stream bool) ollamaChatRequest {
+	msgs := make([]ollamaMsg, len(req.Messages))
+	for i, m := range req.Messages {
+		content, _ := m.Content.(string)
+		msgs[i] = ollamaMsg{Role: m.Role, Content: content}
+	}
+
+	ollamaReq := ollamaChatRequest{
+		Model:    req.Model,
+		Messages: msgs,
+		Stream:   stream,
+	}
+
+	if req.MaxTokens > 0 {
+		ollamaReq.Options.NumPredict = req.MaxTokens
+	}
+	if req.Temperature != nil {
+		ollamaReq.Options.Temperature = *req.Temperature
+	}
+	if req.TopP != nil {
+		ollamaReq.Options.TopP = *req.TopP
+	}
+
+	return ollamaReq
+}