File size: 2,472 Bytes
90a206f
 
621b880
90a206f
621b880
 
2a362b9
621b880
 
 
 
 
 
 
 
 
 
 
 
90a206f
621b880
 
 
 
 
 
074e2f3
 
2a362b9
074e2f3
90a206f
074e2f3
2a362b9
074e2f3
 
 
 
2a362b9
 
074e2f3
 
 
90a206f
a544611
eabfbb2
 
a544611
 
 
90a206f
074e2f3
 
90a206f
 
074e2f3
 
 
 
90a206f
074e2f3
 
 
90a206f
 
 
eabfbb2
074e2f3
90a206f
eabfbb2
 
e1c3d60
 
 
 
 
 
eabfbb2
621b880
f30612d
621b880
eabfbb2
90a206f
621b880
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
// @ts-nocheck
// Preventing TS checks with files presented in the video for a better presentation.
import { streamText as _streamText, convertToCoreMessages } from 'ai';
import { getModel } from '~/lib/.server/llm/model';
import { MAX_TOKENS } from './constants';
import { getSystemPrompt } from './prompts';
import { MODEL_LIST, DEFAULT_MODEL, DEFAULT_PROVIDER, MODEL_REGEX, PROVIDER_REGEX } from '~/utils/constants';

interface ToolResult<Name extends string, Args, Result> {
  toolCallId: string;
  toolName: Name;
  args: Args;
  result: Result;
}

interface Message {
  role: 'user' | 'assistant';
  content: string;
  toolInvocations?: ToolResult<string, unknown, unknown>[];
  model?: string;
}

export type Messages = Message[];

export type StreamingOptions = Omit<Parameters<typeof _streamText>[0], 'model'>;

function extractPropertiesFromMessage(message: Message): { model: string; provider: string; content: string } {
  // Extract model
  const modelMatch = message.content.match(MODEL_REGEX);
  const model = modelMatch ? modelMatch[1] : DEFAULT_MODEL;

  // Extract provider
  const providerMatch = message.content.match(PROVIDER_REGEX);
  const provider = providerMatch ? providerMatch[1] : DEFAULT_PROVIDER;

  // Remove model and provider lines from content
  const cleanedContent = message.content
    .replace(MODEL_REGEX, '')
    .replace(PROVIDER_REGEX, '')
    .trim();

  return { model, provider, content: cleanedContent };
}
export function streamText(
  messages: Messages,
  env: Env,
  options?: StreamingOptions,
  apiKeys?: Record<string, string>
) {
  let currentModel = DEFAULT_MODEL;
  let currentProvider = DEFAULT_PROVIDER;

  const processedMessages = messages.map((message) => {
    if (message.role === 'user') {
      const { model, provider, content } = extractPropertiesFromMessage(message);

      if (MODEL_LIST.find((m) => m.name === model)) {
        currentModel = model;
      }

      currentProvider = provider;

      return { ...message, content };
    }

    return message; 
  });

  const modelDetails = MODEL_LIST.find((m) => m.name === currentModel);

  

    const dynamicMaxTokens =
modelDetails && modelDetails.maxTokenAllowed
  ? modelDetails.maxTokenAllowed
  : MAX_TOKENS;

  return _streamText({
    model: getModel(currentProvider, currentModel, env, apiKeys),
    system: getSystemPrompt(),
    maxTokens: dynamicMaxTokens,
    messages: convertToCoreMessages(processedMessages),
    ...options,
  });
}