247 lines
6.6 KiB
Go
247 lines
6.6 KiB
Go
package openapi
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"fmt"
|
|
"time"
|
|
)
|
|
|
|
// SessionHooks provides hooks into the session lifecycle.
|
|
type SessionHooks struct {
|
|
BeforeMessage func(ctx context.Context, s *Session, msg *Message) error
|
|
AfterMessage func(ctx context.Context, s *Session, msg *Message) error
|
|
BeforeToolCall func(ctx context.Context, s *Session, tc *ToolCall) error
|
|
AfterToolCall func(ctx context.Context, s *Session, tc *ToolCall, result any, err error) error
|
|
OnReasoning func(ctx context.Context, s *Session, reasoning string) error
|
|
OnDone func(ctx context.Context, s *Session) error
|
|
OnError func(ctx context.Context, s *Session, err error)
|
|
}
|
|
|
|
// Session manages a conversation history and tool execution logic.
|
|
type Session struct {
|
|
client *Client
|
|
model string
|
|
messages []Message
|
|
registry ToolRegistry
|
|
hooks SessionHooks
|
|
|
|
// TurnTimeout is the timeout for a single LLM request.
|
|
TurnTimeout time.Duration
|
|
|
|
// DefaultRequestOptions allows setting default parameters for all requests in the session.
|
|
DefaultRequestOptions ChatCompletionRequest
|
|
}
|
|
|
|
// NewSession creates a new session with the given client and model.
|
|
func NewSession(client *Client, model string) *Session {
|
|
return &Session{
|
|
client: client,
|
|
model: model,
|
|
registry: NewMapRegistry(),
|
|
TurnTimeout: 2 * time.Minute,
|
|
DefaultRequestOptions: ChatCompletionRequest{
|
|
CachePrompt: true,
|
|
},
|
|
}
|
|
}
|
|
|
|
// SetRegistry replaces the tool registry.
|
|
func (s *Session) SetRegistry(r ToolRegistry) {
|
|
s.registry = r
|
|
}
|
|
|
|
// RegisterTool is a helper for the default MapRegistry.
|
|
func (s *Session) RegisterTool(name, description string, parameters any, handler ToolHandler) error {
|
|
if mr, ok := s.registry.(*MapRegistry); ok {
|
|
return mr.Register(name, description, parameters, handler)
|
|
}
|
|
return fmt.Errorf("current registry does not support direct registration")
|
|
}
|
|
|
|
// AddMessage adds a message to the conversation history, triggering hooks.
|
|
func (s *Session) AddMessage(ctx context.Context, msg Message) error {
|
|
if s.hooks.BeforeMessage != nil {
|
|
if err := s.hooks.BeforeMessage(ctx, s, &msg); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
s.messages = append(s.messages, msg)
|
|
if s.hooks.AfterMessage != nil {
|
|
return s.hooks.AfterMessage(ctx, s, &msg)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// SetSystemPrompt sets the initial system prompt.
|
|
func (s *Session) SetSystemPrompt(prompt string) {
|
|
if len(s.messages) > 0 && s.messages[0].Role == RoleSystem {
|
|
s.messages[0].Content = prompt
|
|
} else {
|
|
s.messages = append([]Message{{Role: RoleSystem, Content: prompt}}, s.messages...)
|
|
}
|
|
}
|
|
|
|
// Chat sends a user prompt and handles the response loop asynchronously.
|
|
func (s *Session) Chat(ctx context.Context, userPrompt string) (<-chan ChatCompletionChunk, <-chan error) {
|
|
chunks := make(chan ChatCompletionChunk)
|
|
errs := make(chan error, 1)
|
|
go s.runLoop(ctx, userPrompt, chunks, errs)
|
|
return chunks, errs
|
|
}
|
|
|
|
func (s *Session) runLoop(ctx context.Context, userPrompt string, chunks chan<- ChatCompletionChunk, errs chan<- error) {
|
|
defer close(chunks)
|
|
defer close(errs)
|
|
|
|
if userPrompt != "" {
|
|
if err := s.AddMessage(ctx, Message{Role: RoleUser, Content: userPrompt}); err != nil {
|
|
s.handleError(ctx, err, errs)
|
|
return
|
|
}
|
|
}
|
|
|
|
for {
|
|
resp, err := s.nextCompletion(ctx)
|
|
if err != nil {
|
|
s.handleError(ctx, err, errs)
|
|
return
|
|
}
|
|
|
|
choice := resp.Choices[0]
|
|
if err := s.processResponse(ctx, choice.Message); err != nil {
|
|
s.handleError(ctx, err, errs)
|
|
return
|
|
}
|
|
|
|
if len(choice.Message.ToolCalls) > 0 {
|
|
if err := s.handleToolCalls(ctx, choice.Message.ToolCalls); err != nil {
|
|
s.handleError(ctx, err, errs)
|
|
return
|
|
}
|
|
continue
|
|
}
|
|
|
|
if s.hooks.OnDone != nil {
|
|
s.hooks.OnDone(ctx, s)
|
|
}
|
|
chunks <- s.asChunk(resp, choice)
|
|
return
|
|
}
|
|
}
|
|
|
|
func (s *Session) nextCompletion(ctx context.Context) (*ChatCompletionResponse, error) {
|
|
turnCtx, cancel := s.extendContext(ctx)
|
|
defer cancel()
|
|
|
|
req := s.DefaultRequestOptions
|
|
req.Model = s.model
|
|
req.Messages = s.messages
|
|
req.Tools = s.registry.Tools(turnCtx)
|
|
return s.client.CreateChatCompletion(turnCtx, req)
|
|
}
|
|
|
|
func (s *Session) processResponse(ctx context.Context, msg Message) error {
|
|
if msg.ReasoningContent != "" && s.hooks.OnReasoning != nil {
|
|
if err := s.hooks.OnReasoning(ctx, s, msg.ReasoningContent); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return s.AddMessage(ctx, msg)
|
|
}
|
|
|
|
func (s *Session) handleToolCalls(ctx context.Context, toolCalls []ToolCall) error {
|
|
for i := range toolCalls {
|
|
tc := &toolCalls[i]
|
|
if s.hooks.BeforeToolCall != nil {
|
|
if err := s.hooks.BeforeToolCall(ctx, s, tc); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
|
|
// Use extended context for tool execution to prevent timeouts on long tasks
|
|
toolCtx, cancel := s.extendContext(ctx)
|
|
result, err := s.registry.Execute(toolCtx, s, tc.Function.Name, json.RawMessage(tc.Function.Arguments))
|
|
cancel()
|
|
|
|
if s.hooks.AfterToolCall != nil {
|
|
if hErr := s.hooks.AfterToolCall(ctx, s, tc, result, err); hErr != nil {
|
|
return hErr
|
|
}
|
|
}
|
|
|
|
content := ""
|
|
if err != nil {
|
|
content = fmt.Sprintf("Error: %v", err)
|
|
} else {
|
|
resJSON, _ := json.Marshal(result)
|
|
content = string(resJSON)
|
|
}
|
|
|
|
if err := s.AddMessage(ctx, Message{
|
|
Role: RoleTool,
|
|
Name: tc.Function.Name,
|
|
ToolCallID: tc.ID,
|
|
Content: content,
|
|
}); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// extendContext creates a context that resets the timeout for a new "turn",
|
|
// ignoring the parent's deadline but respecting its cancellation.
|
|
func (s *Session) extendContext(ctx context.Context) (context.Context, context.CancelFunc) {
|
|
extCtx, cancel := context.WithTimeout(context.Background(), s.TurnTimeout)
|
|
|
|
stop := make(chan struct{})
|
|
go func() {
|
|
select {
|
|
case <-ctx.Done():
|
|
cancel()
|
|
case <-stop:
|
|
case <-extCtx.Done():
|
|
}
|
|
}()
|
|
|
|
return extCtx, func() {
|
|
close(stop)
|
|
cancel()
|
|
}
|
|
}
|
|
|
|
func (s *Session) asChunk(resp *ChatCompletionResponse, choice Choice) ChatCompletionChunk {
|
|
return ChatCompletionChunk{
|
|
ID: resp.ID,
|
|
Object: "chat.completion.chunk",
|
|
Created: resp.Created,
|
|
Model: resp.Model,
|
|
Choices: []struct {
|
|
Index int `json:"index"`
|
|
Delta Message `json:"delta"`
|
|
FinishReason string `json:"finish_reason"`
|
|
}{
|
|
{
|
|
Index: choice.Index,
|
|
Delta: Message{
|
|
Role: choice.Message.Role,
|
|
Content: choice.Message.Content,
|
|
ReasoningContent: choice.Message.ReasoningContent,
|
|
},
|
|
FinishReason: choice.FinishReason,
|
|
},
|
|
},
|
|
}
|
|
}
|
|
|
|
func (s *Session) handleError(ctx context.Context, err error, errs chan<- error) {
|
|
if s.hooks.OnError != nil {
|
|
s.hooks.OnError(ctx, s, err)
|
|
}
|
|
errs <- err
|
|
}
|
|
|
|
func (s *Session) Messages() []Message { return s.messages }
|
|
func (s *Session) SetHooks(h SessionHooks) { s.hooks = h }
|