diff --git a/main.go b/main.go index 06b606f..81e133a 100644 --- a/main.go +++ b/main.go @@ -13,16 +13,32 @@ import ( "github.com/tmc/langchaingo/memory" ) +const ( + model = "Qwen/Qwen3-Coder-30B-A3B-Instruct" + base_url = "https://api.siliconflow.cn/v1" + token = "sk-udqgogvqgfriqvsehkpxgaejfclvoikbntjeaeijhqnhnviw" + cotPrompt = `请按照以下步骤分析这道算法题: + 1. 算法判断:判断题目类型和所需算法 + 2. 基本思路:说明解题思路和关键点 + 3. 框架代码:给出代码框架 + 4. 易错解析:指出常见错误和注意事项 + + 题目:%s + + 请按上述格式进行分析:` +) + func main() { llm, err := openai.New( - openai.WithModel("Qwen/Qwen3-Coder-30B-A3B-Instruct"), - openai.WithBaseURL("https://api.siliconflow.cn/v1"), - openai.WithToken("sk-udqgogvqgfriqvsehkpxgaejfclvoikbntjeaeijhqnhnviw"), + openai.WithModel(model), + openai.WithBaseURL(base_url), + openai.WithToken(token), ) if err != nil { log.Fatal(err) } + // [Memory] Init chatMemory := memory.NewConversationBuffer() ctx := context.Background() reader := bufio.NewReader(os.Stdin) @@ -39,16 +55,19 @@ func main() { break } - messages, _ := chatMemory.ChatHistory.Messages(ctx) - - // Conversation - var conversation string - for _, msg := range messages { - conversation += msg.GetContent() + "\n" + var fullPrompt string + if isAlgorithmQuestion(userInput) { + fullPrompt = fmt.Sprintf(cotPrompt, userInput) + } else { + messages, _ := chatMemory.ChatHistory.Messages(ctx) + // Conversation + var conversation string + for _, msg := range messages { + conversation += msg.GetContent() + "\n" + } + fullPrompt = conversation + "Human: " + userInput + "\nAssistant:" } - fullPrompt := conversation + "Human: " + userInput + "\nAssistant:" - response, err := llms.GenerateFromSinglePrompt( ctx, llm, @@ -66,3 +85,13 @@ func main() { } } + +func isAlgorithmQuestion(input string) bool { + keywords := []string{"算法", "编程", "题目", "解题", "代码实现"} + for _, keyword := range keywords { + if strings.Contains(input, keyword) { + return true + } + } + return false +}