Atopos

1
import dotenv from "dotenv";
2
dotenv.config();
3

4
export const CONFIG = {
5
  port: parseInt(process.env.WS_PORT || "8080", 10),
6

7
  llm: {
8
    model: "GLM-4.7",
9
    apiKey: process.env.ARK_API_KEY,        // 敏感信息从环境变量读
10
    temperature: 0.7,
11
    timeout: 30000,
12
    maxRetries: 2,
13
    baseURL: "https://ark.cn-beijing.volces.com/api/coding/v3",
14
  },
15

16
  rateLimit: {
17
    windowMs: 60_000,      // 时间窗口：1 分钟
18
    maxRequests: 20,       // 最多 20 次请求
19
  },
20

21
  systemPrompt: "你是一个智能助手，请用友好、自然的方式回答用户的问题。",
22
};

1
import { ChatOpenAI } from "@langchain/openai";
2
import { CONFIG } from "./config";
3

4
/** LLM 实例：纯流式/非流式调用，不绑定工具 */
5
export const llm = new ChatOpenAI({
6
  model: CONFIG.llm.model,
7
  apiKey: CONFIG.llm.apiKey,
8
  temperature: CONFIG.llm.temperature,
9
  streamUsage: false,
10
  timeout: CONFIG.llm.timeout,
11
  maxRetries: CONFIG.llm.maxRetries,
12
  configuration: {
13
    baseURL: CONFIG.llm.baseURL,
14
  },
15
});

1
import { SystemMessage, HumanMessage, AIMessage } from "@langchain/core/messages";
2
import { CONFIG } from "./config";
3

4
export type MemoryStrategy = "buffer" | "window" | "summary" | "vector";
5

6
export interface Session {
7
  id: string;
8
  messages: BaseMessage[];
9
  createdAt: number;
10
  messageCount: number;
11
  // 各种配置都挂在 session 上
12
  memoryStrategy: MemoryStrategy;
13
  windowSize: number;
14
  summaryThreshold: number;
15
  recentToKeep: number;
16
  vectorEntries: VectorEntry[];
17
  currentTemplate?: string;
18
}
19

20
// 内存存储会话（生产环境可以换成 Redis）
21
const sessions = new Map<string, Session>();
22

23
export function createSession(id: string): Session { /* ... */ }
24
export function getSession(id: string): Session | undefined { /* ... */ }
25
export function deleteSession(id: string): void { /* ... */ }
26
export function addUserMessage(session: Session, content: string): void { /* ... */ }
27
export function addAIMessage(session: Session, content: string): void { /* ... */ }
28
export function clearHistory(session: Session): void { /* ... */ }

1
import { PromptTemplate } from "@langchain/core/prompts";
2

3
export const promptTemplates: Record<string, PromptTemplate> = {
4
  tech: PromptTemplate.fromTemplate(`
5
你是一位专业的技术导师。请用简洁易懂的语言回答问题。
6
用户问题：{question}
7
`),
8
  life: PromptTemplate.fromTemplate(`...`),
9
  english: PromptTemplate.fromTemplate(`...`),
10
  interview: PromptTemplate.fromTemplate(`...`),
11
  cot: PromptTemplate.fromTemplate(`...`),
12
  json: PromptTemplate.fromTemplate(`...`),
13
};
14

15
export const DEFAULT_TEMPLATE = "tech";
16

17
export function getPromptTemplate(key?: string): PromptTemplate {
18
  return promptTemplates[key ?? DEFAULT_TEMPLATE] ?? promptTemplates[DEFAULT_TEMPLATE];
19
}

1
import { llm } from "../llm";
2
import { addAIMessage, addUserMessage } from "../session";
3
import { getContextMessages } from "../memory-strategies";
4
import { getPromptTemplate } from "../prompts";
5

6
export async function handleStream(
7
  ws: WebSocket,
8
  session: Session,
9
  content: string,
10
  templateKey: string = "tech",
11
): Promise<void> {
12
  // 1. 参数校验
13
  if (!content.trim()) {
14
    ws.send(JSON.stringify({ type: "error", content: "消息不能为空" }));
15
    return;
16
  }
17

18
  // 2. 用 Prompt 模板包装
19
  const template = getPromptTemplate(templateKey);
20
  const formatted = await template.format({ question: content });
21
  addUserMessage(session, formatted);
22

23
  try {
24
    // 3. 根据记忆策略构建上下文
25
    const contextMessages = await getContextMessages(session, content);
26

27
    // 4. 流式调用 LLM
28
    const stream = await llm.stream(contextMessages);
29
    let collected = "";
30

31
    for await (const chunk of stream) {
32
      const text = typeof chunk.content === "string" ? chunk.content : "";
33
      if (!text) continue;
34
      collected += text;
35
      // 逐字符发送 → 打字机效果
36
      for (const char of text) {
37
        ws.send(JSON.stringify({ type: "chunk", content: char }));
38
      }
39
    }
40

41
    const reply = collected || "（空回复）";
42
    addAIMessage(session, reply);
43

44
    // 5. 向量策略：将对话存入向量库
45
    if (session.memoryStrategy === "vector") {
46
      addToVectorStore(session, content, reply);
47
    }
48

49
    ws.send(JSON.stringify({ type: "done", mode: "stream", content: reply }));
50
  } catch (error: any) {
51
    logger.error(`[${session.id}] stream 调用失败`, error.message);
52
    ws.send(JSON.stringify({ type: "error", content: `请求失败: ${error.message}` }));
53
  }
54
}

1
export async function getContextMessages(
2
  session: Session,
3
  currentInput: string
4
): Promise<BaseMessage[]> {
5
  switch (session.memoryStrategy) {
6
    case "buffer":
7
      return bufferContext(session);
8
    case "window":
9
      return windowContext(session);
10
    case "summary":
11
      return await summaryContext(session);
12
    case "vector":
13
      return await vectorContext(session, currentInput);
14
    default:
15
      return bufferContext(session);
16
  }
17
}
18

19
// ---------- 四种策略的具体实现 ----------
20

21
function bufferContext(session: Session): BaseMessage[] {
22
  // 全量保留所有消息
23
  return session.messages;
24
}
25

26
function windowContext(session: Session): BaseMessage[] {
27
  // 只保留最近 N 轮对话
28
  const system = session.messages.find(m => m._getType() === "system");
29
  const recent = session.messages.slice(-session.windowSize * 2);
30
  return system ? [system, ...recent] : recent;
31
}
32

33
async function summaryContext(session: Session): Promise<BaseMessage[]> {
34
  // 超过阈值触发压缩，旧消息变摘要，只保留最近几条
35
  if (nonSystem.length >= session.summaryThreshold) {
36
    session.summary = await compressToSummary(...);
37
  }
38
  return buildContextFromSummary(session);
39
}
40

41
async function vectorContext(session: Session, input: string): Promise<BaseMessage[]> {
42
  // 向量检索相关历史 + 最近对话
43
  const related = findRelatedHistory(session.vectorEntries, input);
44
  const recent = session.messages.slice(-4);
45
  return [...related, ...recent];
46
}

1
import { CONFIG } from "./config";
2

3
interface RateLimitEntry {
4
  count: number;
5
  resetTime: number;
6
}
7

8
const rateLimitMap = new Map<string, RateLimitEntry>();
9

10
export function checkRateLimit(sessionId: string): boolean {
11
  const now = Date.now();
12
  const entry = rateLimitMap.get(sessionId);
13

14
  // 窗口过期，重置计数
15
  if (!entry || now >= entry.resetTime) {
16
    rateLimitMap.set(sessionId, {
17
      count: 1,
18
      resetTime: now + CONFIG.rateLimit.windowMs,
19
    });
20
    return true;
21
  }
22

23
  // 在窗口内，检查是否超限
24
  if (entry.count >= CONFIG.rateLimit.maxRequests) {
25
    return false;
26
  }
27

28
  entry.count++;
29
  return true;
30
}
31

32
export function clearRateLimit(sessionId: string): void {
33
  rateLimitMap.delete(sessionId);
34
}