const tools = [{
  type: "function",
  function: {
    name: "queryPortKB",
    description: "检索码头内部知识库，包括操作手册、安全规范、设备参数。适用于回答具体的业务流程或技术参数问题。",
    parameters: {
      type: "object",
      properties: {
        query: { 
          type: "string", 
          description: "经过重写的、包含专业术语的搜索查询语句。例如将'吊车坏了'重写为'起重机故障排查'" 
        },
        category: { 
          type: "string", 
          enum: ["safety", "operation", "maintenance"], 
          description: "可选的知识分类，用于缩小检索范围" 
        }
      },
      required: ["query"]
    }
  }
}];

while (iteration < max_iterations) {
  const response = await llm.chat({ messages: history, tools });
  
  if (response.tool_calls) {
    // 执行专属工具
    const result = await queryPortKB(response.tool_calls[0].arguments.query);
    // 回传结果 (关键：带上 tool_call_id)
    history.push({ role: "tool", content: result, tool_call_id: response.tool_calls[0].id });
  } else {
    // 模型已汇总知识，输出最终答案
    return response.content; 
  }
}

-- 1. 启用 pgvector 插件
CREATE EXTENSION IF NOT EXISTS vector;

-- 2. 创建表
CREATE TABLE port_knowledge_base (
    id BIGSERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536),
    doc_type VARCHAR(50),
    source_file VARCHAR(255),
    page_num INTEGER,
    security_level VARCHAR(20) DEFAULT 'public',
    created_at TIMESTAMP DEFAULT NOW()
);

-- 3. 创建索引
CREATE INDEX ON port_knowledge_base USING ivfflat (embedding vector_cosine_ops) WITH (lists = 100);
CREATE INDEX idx_doc_type ON port_knowledge_base(doc_type);

SELECT 
    content, 
    source_file, 
    page_num, 
    1 - (embedding <=> $1) AS similarity_score -- 计算相似度 (0~1)
FROM port_knowledge_base
WHERE 
    doc_type = ANY($2)          -- 过滤类型: 比如只查 ['safety', 'operation']
    AND security_level = $3     -- 过滤权限: 比如当前用户只能查 'public'
ORDER BY 
    embedding <=> $1            -- 按相似度排序 (越近越小)
LIMIT 5;                        -- 取前 5 个最相关的片段

interface Chunk {
  content: string;
  metadata: {
    sourceFile: string;
    pageNum: number;
    docType: 'safety' | 'operation' | 'maintenance';
  };
}

function smartChunk(text: string, fileName: string): Chunk[] {
  const chunks: Chunk[] = [];
  
  // 1. 按章节标题分割 (Regex 匹配 "## 第 X 章" 或 "1.2 安全规范")
  const sections = text.split(/(?=^##\s|^\d+\.\d+\s)/m); 

  sections.forEach(section => {
    // 2. 如果章节太长 (>1000 字)，再按段落二次分割
    if (section.length > 1000) {
      const paragraphs = section.split('\n\n');
      let currentChunk = "";
      
      paragraphs.forEach(p => {
        if ((currentChunk + p).length > 800) {
          // 保存当前块
          chunks.push({ 
            content: currentChunk, 
            metadata: { sourceFile: fileName, pageNum: extractPage(section), docType: inferType(section) } 
          });
          // 保留 20% 重叠
          currentChunk = currentChunk.slice(-160) + p; 
        } else {
          currentChunk += p + "\n\n";
        }
      });
      // 收尾
      if (currentChunk.trim()) {
        chunks.push({ 
          content: currentChunk, 
          metadata: { sourceFile: fileName, pageNum: extractPage(section), docType: inferType(section) } 
        });
      }
    } else {
      // 短章节直接存
      chunks.push({ 
        content: section, 
        metadata: { sourceFile: fileName, pageNum: extractPage(section), docType: inferType(section) } 
      });
    }
  });
  
  return chunks;
}

mkdir port-agent-project && cd port-agent-project
mkdir server && cd server
npm init -y
npm install express openai pg dotenv cors
npm install -D typescript @types/node @types/express @types/pg nodemon ts-node
npx tsc --init

# OpenAI API (或其他兼容模型)
OPENAI_API_KEY=sk-...

# PostgreSQL 连接串 (包含 pgvector)
# 格式: postgres://用户:密码@主机:端口/数据库名
DATABASE_URL=postgres://postgres:password@localhost:5432/port_agent_db

# 安全配置
PORT=3001

/*
 * @Author: 刘书记
 * @Date: 2026-04-01 17:18:00
 * @Description: 努力搬砖的小花花
 */
import { Pool } from 'pg';
import OpenAI from 'openai';
import dotenv from 'dotenv';

dotenv.config();

// 初始化客户端
const pool = new Pool({ connectionString: process.env.DATABASE_URL });
const openai = new OpenAI({ 
  apiKey: process.env.OPENAI_API_KEY,
  baseURL: process.env.OPENAI_BASE_URL // 支持中转 API
});

// 1. 定义工具的描述 (用于 Function Calling)
export const queryPortKBTool = {
  type: "function" as const,
  function: {
    name: "queryPortKB",
    description: "检索码头内部知识库，包括操作手册、安全规范。适用于回答具体的业务流程或技术参数问题。",
    parameters: {
      type: "object",
      properties: {
        query: { 
          type: "string", 
          description: "经过重写的、包含专业术语的搜索查询语句。例如将'吊车坏了'重写为'起重机故障排查'" 
        },
        docType: { 
          type: "string", 
          enum: ["safety", "operation", "maintenance"], 
          description: "可选的文档类型过滤" 
        }
      },
      required: ["query"]
    }
  }
};

// 2. 真正的执行函数
export async function executeQueryPortKB(args: { query: string; docType?: string }) {
  const { query, docType } = args;
  console.log(`🔍 正在检索知识库: "${query}" (类型: ${docType || '全部'})...`);

  try {
    // A. 生成查询向量
    const embeddingResponse = await openai.embeddings.create({
      model: "text-embedding-3-small",
      input: query,
    });
    const vector = embeddingResponse?.data?.[0]?.embedding;
    if (!vector || vector.length === 0) {
      console.log("⚠️ 向量生成结果为空，跳过知识库检索");
      return "当前问题未能生成有效检索向量，请稍后重试或换个问法。";
    }

    // B. 构建 SQL (带元数据过滤)
    // 注意：$1 是向量，$2 是 docType (如果有)
    let sql = `
      SELECT content, source_file, page_num, 1 - (embedding <=> $1::vector) AS similarity
      FROM port_knowledge_base
      WHERE 1=1
    `;
    const vectorLiteral = `[${vector.join(",")}]`;
    const values: any[] = [vectorLiteral];
    
    if (docType) {
      sql += ` AND doc_type = $2`;
      values.push(docType);
    }

    // 按相似度排序，取前 5 个
    sql += ` ORDER BY embedding <=> $1::vector LIMIT 5`;

    const res = await pool.query(sql, values);

    // C. ⚠️ 关键防御：相似度阈值检查
    if (res.rows.length === 0 || res.rows[0].similarity < 0.4) {
      console.log('⚠️ 未找到高相关性的文档 (最高分:', res.rows[0]?.similarity?.toFixed(3) || 0, ')');
      return "未在码头知识库中找到与当前问题高度相关的文档。请尝试换个问法或确认问题是否属于业务范围。";
    }

    // D. 格式化返回结果 (带引用信息)
    const context = res.rows.map((row, idx) => 
      `[片段 ${idx + 1} 来源：${row.source_file} 第${row.page_num}页 | 相似度: ${row.similarity.toFixed(3)}]\n${row.content}`
    ).join('\n\n');

    console.log(`✅ 找到 ${res.rows.length} 个相关片段，最高相似度: ${res.rows[0].similarity.toFixed(3)}`);
    return context;

  } catch (error) {
    console.error("❌ DB Query Error:", error);
    throw new Error(`知识库检索失败: ${error instanceof Error ? error.message : '未知错误'}`);
  }
}

import { executeQueryPortKB } from './agent/tools';

async function testTool() {
  try {
    // 模拟一个用户查询 (假设你库里已经有数据)
    // 如果库是空的，这个测试会返回 "未找到..."，那也是正常的
    const result = await executeQueryPortKB({ 
      query: "起重机 紧急制动 操作流程", 
      docType: "operation" 
    });
    
    console.log("\n--- 工具返回结果 ---");
    console.log(result);
  } catch (err) {
    console.error(err);
  }
}

testTool();

npx ts-node src/test-tool.ts

import { Pool } from 'pg';
import OpenAI from 'openai';
import dotenv from 'dotenv';

dotenv.config();

const pool = new Pool({ connectionString: process.env.DATABASE_URL });
const openai = new OpenAI({ 
  apiKey: process.env.OPENAI_API_KEY,
  baseURL: process.env.OPENAI_BASE_URL 
});

const testData = [
  {
    content: "紧急制动操作流程：1. 立即按下红色紧急停止按钮。2. 将主控制器手柄归零。3. 通过对讲机通知控制室。4. 等待维修人员检查液压系统压力。",
    doc_type: "operation",
    source_file: "起重机操作手册_v2.3.pdf",
    page_num: 42
  },
  {
    content: "安全规范第 5 条：所有操作人员必须佩戴安全帽和反光背心。严禁在吊臂下方站立或通行。风速超过 12m/s 时必须停止作业。",
    doc_type: "safety",
    source_file: "港口安全红宝书_2024.pdf",
    page_num: 15
  },
  {
    content: "日常维护检查表：每日开班前需检查钢丝绳磨损情况、制动器灵敏度、限位开关是否正常工作。发现异常立即上报。",
    doc_type: "maintenance",
    source_file: "设备维护保养指南.pdf",
    page_num: 8
  },
  {
    content: "故障代码 E-03 表示液压系统压力过低。处理方法：检查油箱液位，补充液压油；检查油泵皮带是否松动。",
    doc_type: "maintenance",
    source_file: "起重机故障代码速查表.pdf",
    page_num: 3
  }
];

async function seed() {
  console.log('🌱 开始插入测试数据...');
  
  for (const item of testData) {
    try {
      // 1. 生成向量
      const embedding = await openai.embeddings.create({
        model: "text-embedding-3-small",
        input: item.content,
      });
      const vector = embedding.data[0].embedding;

      // 2. 插入数据库
      await pool.query(
        `INSERT INTO port_knowledge_base (content, embedding, doc_type, source_file, page_num)
         VALUES ($1, $2, $3, $4, $5)`,
        [item.content, vector, item.doc_type, item.source_file, item.page_num]
      );
      console.log(`✅ 已插入: "${item.content.substring(0, 20)}..." (${item.doc_type})`);
    } catch (err) {
      console.error(`❌ 插入失败:`, err);
    }
  }

  console.log('\n🎉 数据种子完成！现在可以运行 test-tool.ts 进行测试了。');
  await pool.end();
}

seed();

npx ts-node src/seed-data.ts

npx ts-node src/test-tool.ts

export const SYSTEM_PROMPT = `
你是一个专业的码头作业助手，专门回答关于起重机操作、安全规范和设备维护的问题。

**核心指令：**
1. **优先检索**：当用户询问具体的操作流程、故障代码、安全规定时，**必须**优先调用 \`queryPortKB\` 工具检索知识库。严禁凭空编造答案。
2. **查询重写**：在调用工具前，请将用户的口语化问题重写为专业的工程术语。
   - 用户问："吊车坏了咋办？" -> 你调用工具时用："起重机 常见故障 排查流程"
   - 用户问："那个红线是啥？" -> 你调用工具时用："安全规范 警戒线 定义"
3. **引用来源**：回答时必须注明信息来源，例如：“根据《起重机操作手册》第 42 页...”。
4. **边界控制**：如果工具返回“未找到相关文档”，请诚实告知用户，不要强行回答。
`;

/*
 * @Author: 刘书记
 * @Date: 2026-04-01 18:26:41
 * @Description: 努力搬砖的小花花
 */
import express from 'express';
import cors from 'cors';
import OpenAI from 'openai';
import dotenv from 'dotenv';
import { queryPortKBTool, executeQueryPortKB } from './agent/tools';
import { SYSTEM_PROMPT } from './agent/systemPrompt';

dotenv.config();

const app = express();
app.use(cors());
app.use(express.json());

const openai = new OpenAI({ 
  apiKey: process.env.OPENAI_API_KEY,
  baseURL: process.env.OPENAI_BASE_URL 
});

app.post('/api/chat', async (req, res) => {
  const userMessage = req.body?.message ?? req.body?.maessage;
  if (typeof userMessage !== 'string' || userMessage.trim().length === 0) {
    return res.status(400).json({ error: "请求体缺少有效的 message 字段（字符串）" });
  }

  // 设置 SSE 响应头
  res.setHeader('Content-Type', 'text/event-stream; charset=utf-8');
  res.setHeader('Cache-Control', 'no-cache');
  res.setHeader('Connection', 'keep-alive');

  const inputHistory = Array.isArray(req.body?.history) ? req.body.history : [];
  const chatHistory: any[] = [...inputHistory, { role: 'user', content: userMessage.trim() }];

  let iteration = 0;
  const maxIterations = 5;
  const sources: Array<{ file: string; page: string }> = [];

  try {
    while (iteration < maxIterations) {
      res.write(`data: ${JSON.stringify({ type: 'thinking', content: '正在分析用户意图...' })}\n\n`);

      const response = await openai.chat.completions.create({
        model: "gpt-4o",
        messages: [{ role: 'system', content: SYSTEM_PROMPT }, ...chatHistory],
        tools: [queryPortKBTool],
        stream: true
      });

      let fullContent = "";
      const toolCallsMap = new Map<number, any>();

      for await (const chunk of response) {
        const delta = chunk.choices[0]?.delta;
        if (!delta) continue;

        if (delta.content) {
          fullContent += delta.content;
          res.write(`data: ${JSON.stringify({ type: 'content', content: delta.content })}\n\n`);
        }

        if (delta.tool_calls) {
          for (const partialToolCall of delta.tool_calls) {
            const idx = partialToolCall.index ?? 0;
            const current = toolCallsMap.get(idx) ?? {
              id: partialToolCall.id ?? "",
              type: 'function',
              function: { name: "", arguments: "" }
            };

            if (partialToolCall.id) current.id = partialToolCall.id;
            if (partialToolCall.function?.name) current.function.name = partialToolCall.function.name;
            if (partialToolCall.function?.arguments) {
              current.function.arguments += partialToolCall.function.arguments;
            }

            toolCallsMap.set(idx, current);
          }
        }
      }

      const toolCalls = Array.from(toolCallsMap.values());

      // 没有工具调用，说明可以直接结束
      if (toolCalls.length === 0) {
        if (sources.length > 0) {
          res.write(`data: ${JSON.stringify({ type: 'sources', data: sources })}\n\n`);
        }
        res.write(`data: [DONE]\n\n`);
        res.end();
        return;
      }

      // 关键：先把 assistant tool_calls 放入历史，再追加 tool 消息
      chatHistory.push({
        role: 'assistant',
        content: fullContent || "",
        tool_calls: toolCalls
      });

      for (const toolCall of toolCalls) {
        if (toolCall.type !== 'function' || toolCall.function?.name !== 'queryPortKB') {
          continue;
        }

        let args: { query: string; docType?: string };
        try {
          args = JSON.parse(toolCall.function.arguments);
        } catch {
          res.write(`data: ${JSON.stringify({ type: 'error', content: '工具参数解析失败，已跳过。' })}\n\n`);
          continue;
        }

        res.write(`data: ${JSON.stringify({
          type: 'tool_start',
          content: `正在检索知识库：${args.query}...`
        })}\n\n`);

        const toolResult = await executeQueryPortKB(args);

        const sourceRegex = /\[片段\s+\d+\s+来源：(.+?)\s+第(\d+)页/g;
        for (const match of toolResult.matchAll(sourceRegex)) {
          const file = match[1];
          const page = match[2];
          if (file && page) {
            sources.push({ file, page });
          }
        }

        chatHistory.push({
          role: 'tool',
          tool_call_id: toolCall.id,
          content: toolResult
        });
      }

      iteration++;
    }

    res.write(`data: ${JSON.stringify({ type: 'error', content: '达到最大推理轮数，已停止。' })}\n\n`);
    res.write(`data: [DONE]\n\n`);
    res.end();
  } catch (error) {
    console.error(error);
    res.write(`data: ${JSON.stringify({ type: 'error', content: '出错了...' })}\n\n`);
    res.end();
  }
});

const PORT = process.env.PORT || 3001;
app.listen(PORT, () => {
  console.log(`🚀 码头 Agent 服务已启动 http://localhost:${PORT}`);
});