mirror of
https://github.com/alibaba/higress.git
synced 2026-03-02 07:30:49 +08:00
Add ai search plugin (#1804)
This commit is contained in:
澄潭
GitHub
@@ -23,7 +23,7 @@ const (
|
||||
SKIP_CACHE_HEADER = "x-higress-skip-ai-cache"
|
||||
ERROR_PARTIAL_MESSAGE_KEY = "errorPartialMessage"
|
||||
|
||||
DEFAULT_MAX_BODY_BYTES uint32 = 10 * 1024 * 1024
|
||||
DEFAULT_MAX_BODY_BYTES uint32 = 100 * 1024 * 1024
|
||||
)
|
||||
|
||||
func main() {
|
||||
|
||||
@@ -20,7 +20,7 @@ import (
|
||||
const (
|
||||
pluginName = "ai-proxy"
|
||||
|
||||
defaultMaxBodyBytes uint32 = 10 * 1024 * 1024
|
||||
defaultMaxBodyBytes uint32 = 100 * 1024 * 1024
|
||||
)
|
||||
|
||||
func main() {
|
||||
|
||||
224
plugins/wasm-go/extensions/ai-search/README.md
Normal file
224
plugins/wasm-go/extensions/ai-search/README.md
Normal file
@@ -0,0 +1,224 @@
|
||||
## 简介
|
||||
---
|
||||
title: AI 搜索增强
|
||||
keywords: [higress,ai search]
|
||||
description: higress 支持通过集成搜索引擎（Google/Bing/Arxiv/Elasticsearch等）的实时结果，增强DeepSeek-R1等模型等回答准确性和时效性
|
||||
---
|
||||
|
||||
## 功能说明
|
||||
|
||||
`ai-search`插件通过集成搜索引擎（Google/Bing/Arxiv/Elasticsearch等）的实时结果，增强AI模型的回答准确性和时效性。插件会自动将搜索结果注入到提示模板中，并根据配置决定是否在最终回答中添加引用来源。
|
||||
|
||||
## 运行属性
|
||||
|
||||
插件执行阶段：`默认阶段`
|
||||
插件执行优先级：`440`
|
||||
|
||||
## 配置字段
|
||||
|
||||
| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
|
||||
|------|----------|----------|--------|------|
|
||||
| needReference | bool | 选填 | false | 是否在回答中添加引用来源 |
|
||||
| referenceFormat | string | 选填 | `"**References:**\n%s"` | 引用内容格式，必须包含%s占位符 |
|
||||
| defaultLang | string | 选填 | - | 默认搜索语言代码（如zh-CN/en-US） |
|
||||
| promptTemplate | string | 选填 | 内置模板 | 提示模板，必须包含`{search_results}`和`{question}`占位符 |
|
||||
| searchFrom | array of object | 必填 | - | 参考下面搜索引擎配置，至少配置一个引擎 |
|
||||
| searchRewrite | object | 选填 | - | 搜索重写配置，用于使用LLM服务优化搜索查询 |
|
||||
|
||||
## 搜索重写说明
|
||||
|
||||
搜索重写功能使用LLM服务对用户的原始查询进行分析和优化，可以：
|
||||
1. 将用户的自然语言查询转换为更适合搜索引擎的关键词组合
|
||||
2. 对于Arxiv论文搜索，自动识别相关的论文类别并添加类别限定
|
||||
3. 对于私有知识库搜索，将长查询拆分成多个精准的关键词组合
|
||||
|
||||
强烈建议在使用Arxiv或Elasticsearch引擎时启用此功能。对于Arxiv搜索，它能准确识别论文所属领域并优化英文关键词；对于私有知识库搜索，它能提供更精准的关键词匹配，显著提升搜索效果。
|
||||
|
||||
## 搜索重写配置
|
||||
|
||||
| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
|
||||
|------|----------|----------|--------|------|
|
||||
| llmServiceName | string | 必填 | - | LLM服务名称 |
|
||||
| llmServicePort | number | 必填 | - | LLM服务端口 |
|
||||
| llmApiKey | string | 必填 | - | LLM服务API密钥 |
|
||||
| llmUrl | string | 必填 | - | LLM服务API地址 |
|
||||
| llmModelName | string | 必填 | - | LLM模型名称 |
|
||||
| timeoutMillisecond | number | 选填 | 30000 | API调用超时时间（毫秒） |
|
||||
|
||||
## 搜索引擎通用配置
|
||||
|
||||
| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
|
||||
|------|----------|----------|--------|------|
|
||||
| type | string | 必填 | - | 引擎类型（google/bing/arxiv/elasticsearch） |
|
||||
| apiKey | string | 必填 | - | 搜索引擎API密钥 |
|
||||
| serviceName | string | 必填 | - | 后端服务名称 |
|
||||
| servicePort | number | 必填 | - | 后端服务端口 |
|
||||
| count | number | 选填 | 10 | 单次搜索返回结果数量 |
|
||||
| start | number | 选填 | 0 | 搜索结果偏移量（从第start+1条结果开始返回） |
|
||||
| timeoutMillisecond | number | 选填 | 5000 | API调用超时时间（毫秒） |
|
||||
| optionArgs | map | 选填 | - | 搜索引擎特定参数（key-value格式） |
|
||||
|
||||
## Google 特定配置
|
||||
|
||||
| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
|
||||
|------|----------|----------|--------|------|
|
||||
| cx | string | 必填 | - | Google自定义搜索引擎ID，用于指定搜索范围 |
|
||||
|
||||
## Arxiv 特定配置
|
||||
|
||||
| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
|
||||
|------|----------|----------|--------|------|
|
||||
| arxivCategory | string | 选填 | - | 搜索的论文[类别](https://arxiv.org/category_taxonomy)（如cs.AI, cs.CL等） |
|
||||
|
||||
## Elasticsearch 特定配置
|
||||
|
||||
| 名称 | 数据类型 | 填写要求 | 默认值 | 描述 |
|
||||
|------|----------|----------|--------|------|
|
||||
| index | string | 必填 | - | 要搜索的Elasticsearch索引名称 |
|
||||
| contentField | string | 必填 | - | 要查询的内容字段名称 |
|
||||
| linkField | string | 必填 | - | 结果链接字段名称 |
|
||||
| titleField | string | 必填 | - | 结果标题字段名称 |
|
||||
|
||||
|
||||
## 配置示例
|
||||
|
||||
### 基础配置（单搜索引擎）
|
||||
|
||||
```yaml
|
||||
needReference: true
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
count: 5
|
||||
optionArgs:
|
||||
fileType: "pdf"
|
||||
|
||||
### Arxiv搜索配置
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: arxiv
|
||||
serviceName: "arxiv-svc.dns"
|
||||
servicePort: 443
|
||||
arxivCategory: "cs.AI"
|
||||
count: 10
|
||||
```
|
||||
|
||||
### 多搜索引擎配置
|
||||
|
||||
```yaml
|
||||
defaultLang: "en-US"
|
||||
promptTemplate: |
|
||||
# Search Results:
|
||||
{search_results}
|
||||
|
||||
# Please answer this question:
|
||||
{question}
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "google-key"
|
||||
cx: "github-search-id" # 专门搜索GitHub内容的搜索引擎ID
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
- type: google
|
||||
apiKey: "google-key"
|
||||
cx: "news-search-id" # 专门搜索Google News内容的搜索引擎ID
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
- type: being
|
||||
apiKey: "bing-key"
|
||||
serviceName: "bing-svc.dns"
|
||||
servicePort: 443
|
||||
optionArgs:
|
||||
answerCount: "5"
|
||||
```
|
||||
|
||||
### 并发查询配置
|
||||
|
||||
由于搜索引擎对单次查询返回结果数量有限制（如Google限制单次最多返回100条结果），可以通过以下方式获取更多结果：
|
||||
1. 设置较小的count值（如10）
|
||||
2. 通过start参数指定结果偏移量
|
||||
3. 并发发起多个查询请求，每个请求的start值按count递增
|
||||
|
||||
例如，要获取30条结果，可以配置count=10并并发发起20个查询，每个查询的start值分别为0,10,20：
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
start: 0
|
||||
count: 10
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
start: 10
|
||||
count: 10
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
start: 20
|
||||
count: 10
|
||||
```
|
||||
|
||||
注意，过高的并发可能会导致限流，需要根据实际情况调整。
|
||||
|
||||
### Elasticsearch 配置（用于对接私有知识库）
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: elasticsearch
|
||||
serviceName: "es-svc.static"
|
||||
# 固定地址服务的端口默认是80
|
||||
servicePort: 80
|
||||
index: "knowledge_base"
|
||||
contentField: "content"
|
||||
linkField: "url"
|
||||
titleField: "title"
|
||||
```
|
||||
|
||||
### 自定义引用格式
|
||||
|
||||
```yaml
|
||||
needReference: true
|
||||
referenceFormat: "### 数据来源\n%s"
|
||||
searchFrom:
|
||||
- type: being
|
||||
apiKey: "your-bing-key"
|
||||
serviceName: "search-service.dns"
|
||||
servicePort: 8080
|
||||
```
|
||||
|
||||
### 搜索重写配置
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
searchRewrite:
|
||||
llmServiceName: "llm-svc.dns"
|
||||
llmServicePort: 443
|
||||
llmApiKey: "your-llm-api-key"
|
||||
llmUrl: "https://api.example.com/v1/chat/completions"
|
||||
llmModelName: "gpt-3.5-turbo"
|
||||
timeoutMillisecond: 15000
|
||||
```
|
||||
|
||||
## 注意事项
|
||||
|
||||
1. 提示词模版必须包含`{search_results}`和`{question}`占位符，可选使用`{cur_date}`插入当前日期（格式：2006年1月2日）
|
||||
2. 默认模板包含搜索结果处理指引和回答规范，如无特殊需要可以直接用默认模板，否则请根据实际情况修改
|
||||
3. 多个搜索引擎是并行查询，总超时时间 = 所有搜索引擎配置中最大timeoutMillisecond值 + 处理时间
|
||||
4. Arxiv搜索不需要API密钥，但可以指定论文类别（arxivCategory）来缩小搜索范围
|
||||
225
plugins/wasm-go/extensions/ai-search/README_EN.md
Normal file
225
plugins/wasm-go/extensions/ai-search/README_EN.md
Normal file
@@ -0,0 +1,225 @@
|
||||
## Introduction
|
||||
---
|
||||
title: AI Search Enhancement
|
||||
keywords: [higress, ai search]
|
||||
description: Higress supports enhancing the accuracy and timeliness of responses from models like DeepSeek-R1 by integrating real-time results from search engines (Google/Bing/Arxiv/Elasticsearch etc.)
|
||||
---
|
||||
|
||||
## Feature Description
|
||||
|
||||
The `ai-search` plugin enhances the accuracy and timeliness of AI model responses by integrating real-time results from search engines (Google/Bing/Arxiv/Elasticsearch etc.). The plugin automatically injects search results into the prompt template and determines whether to add reference sources in the final response based on configuration.
|
||||
|
||||
## Runtime Properties
|
||||
|
||||
Plugin execution stage: `Default stage`
|
||||
Plugin execution priority: `440`
|
||||
|
||||
## Configuration Fields
|
||||
|
||||
| Name | Data Type | Requirement | Default Value | Description |
|
||||
|------|-----------|-------------|---------------|-------------|
|
||||
| needReference | bool | Optional | false | Whether to add reference sources in the response |
|
||||
| referenceFormat | string | Optional | `"**References:**\n%s"` | Reference content format, must include %s placeholder |
|
||||
| defaultLang | string | Optional | - | Default search language code (e.g. zh-CN/en-US) |
|
||||
| promptTemplate | string | Optional | Built-in template | Prompt template, must include `{search_results}` and `{question}` placeholders |
|
||||
| searchFrom | array of object | Required | - | Refer to search engine configuration below, at least one engine must be configured |
|
||||
| searchRewrite | object | Optional | - | Search rewrite configuration, used to optimize search queries using an LLM service |
|
||||
|
||||
## Search Rewrite Description
|
||||
|
||||
The search rewrite feature uses an LLM service to analyze and optimize the user's original query, which can:
|
||||
1. Convert natural language queries into keyword combinations better suited for search engines
|
||||
2. For Arxiv paper searches, automatically identify relevant paper categories and add category constraints
|
||||
3. For private knowledge base searches, break down long queries into multiple precise keyword combinations
|
||||
|
||||
It is strongly recommended to enable this feature when using Arxiv or Elasticsearch engines. For Arxiv searches, it can accurately identify paper domains and optimize English keywords; for private knowledge base searches, it can provide more precise keyword matching, significantly improving search effectiveness.
|
||||
|
||||
## Search Rewrite Configuration
|
||||
|
||||
| Name | Data Type | Requirement | Default Value | Description |
|
||||
|------|-----------|-------------|---------------|-------------|
|
||||
| llmServiceName | string | Required | - | LLM service name |
|
||||
| llmServicePort | number | Required | - | LLM service port |
|
||||
| llmApiKey | string | Required | - | LLM service API key |
|
||||
| llmUrl | string | Required | - | LLM service API URL |
|
||||
| llmModelName | string | Required | - | LLM model name |
|
||||
| timeoutMillisecond | number | Optional | 30000 | API call timeout (milliseconds) |
|
||||
|
||||
## Search Engine Common Configuration
|
||||
|
||||
| Name | Data Type | Requirement | Default Value | Description |
|
||||
|------|-----------|-------------|---------------|-------------|
|
||||
| type | string | Required | - | Engine type (google/bing/arxiv/elasticsearch) |
|
||||
| apiKey | string | Required | - | Search engine API key |
|
||||
| serviceName | string | Required | - | Backend service name |
|
||||
| servicePort | number | Required | - | Backend service port |
|
||||
| count | number | Optional | 10 | Number of results returned per search |
|
||||
| start | number | Optional | 0 | Search result offset (start returning from the start+1 result) |
|
||||
| timeoutMillisecond | number | Optional | 5000 | API call timeout (milliseconds) |
|
||||
| optionArgs | map | Optional | - | Search engine specific parameters (key-value format) |
|
||||
|
||||
## Google Specific Configuration
|
||||
|
||||
| Name | Data Type | Requirement | Default Value | Description |
|
||||
|------|-----------|-------------|---------------|-------------|
|
||||
| cx | string | Required | - | Google Custom Search Engine ID, used to specify search scope |
|
||||
|
||||
## Arxiv Specific Configuration
|
||||
|
||||
| Name | Data Type | Requirement | Default Value | Description |
|
||||
|------|-----------|-------------|---------------|-------------|
|
||||
| arxivCategory | string | Optional | - | Search paper [category](https://arxiv.org/category_taxonomy) (e.g. cs.AI, cs.CL etc.) |
|
||||
|
||||
## Elasticsearch Specific Configuration
|
||||
|
||||
| Name | Data Type | Requirement | Default Value | Description |
|
||||
|------|-----------|-------------|---------------|-------------|
|
||||
| index | string | Required | - | Elasticsearch index name to search |
|
||||
| contentField | string | Required | - | Content field name to query |
|
||||
| linkField | string | Required | - | Result link field name |
|
||||
| titleField | string | Required | - | Result title field name |
|
||||
|
||||
|
||||
## Configuration Examples
|
||||
|
||||
### Basic Configuration (Single Search Engine)
|
||||
|
||||
```yaml
|
||||
needReference: true
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
count: 5
|
||||
optionArgs:
|
||||
fileType: "pdf"
|
||||
```
|
||||
|
||||
### Arxiv Search Configuration
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: arxiv
|
||||
serviceName: "arxiv-svc.dns"
|
||||
servicePort: 443
|
||||
arxivCategory: "cs.AI"
|
||||
count: 10
|
||||
```
|
||||
|
||||
### Multiple Search Engines Configuration
|
||||
|
||||
```yaml
|
||||
defaultLang: "en-US"
|
||||
promptTemplate: |
|
||||
# Search Results:
|
||||
{search_results}
|
||||
|
||||
# Please answer this question:
|
||||
{question}
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "google-key"
|
||||
cx: "github-search-id" # Search engine ID specifically for GitHub content
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
- type: google
|
||||
apiKey: "google-key"
|
||||
cx: "news-search-id" # Search engine ID specifically for Google News content
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
- type: bing
|
||||
apiKey: "bing-key"
|
||||
serviceName: "bing-svc.dns"
|
||||
servicePort: 443
|
||||
optionArgs:
|
||||
answerCount: "5"
|
||||
```
|
||||
|
||||
### Concurrent Query Configuration
|
||||
|
||||
Since search engines limit the number of results per query (e.g. Google limits to 100 results per query), you can get more results by:
|
||||
1. Setting a smaller count value (e.g. 10)
|
||||
2. Specifying result offset with start parameter
|
||||
3. Concurrently initiating multiple query requests, with each request's start value incrementing by count
|
||||
|
||||
For example, to get 30 results, configure count=10 and concurrently initiate 3 queries with start values 0,10,20 respectively:
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
start: 0
|
||||
count: 10
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
start: 10
|
||||
count: 10
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
start: 20
|
||||
count: 10
|
||||
```
|
||||
|
||||
Note that excessive concurrency may lead to rate limiting, adjust according to actual situation.
|
||||
|
||||
### Elasticsearch Configuration (For Private Knowledge Base Integration)
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: elasticsearch
|
||||
serviceName: "es-svc.static"
|
||||
# static ip service use 80 as default port
|
||||
servicePort: 80
|
||||
index: "knowledge_base"
|
||||
contentField: "content"
|
||||
linkField: "url"
|
||||
titleField: "title"
|
||||
```
|
||||
|
||||
### Custom Reference Format
|
||||
|
||||
```yaml
|
||||
needReference: true
|
||||
referenceFormat: "### Data Sources\n%s"
|
||||
searchFrom:
|
||||
- type: bing
|
||||
apiKey: "your-bing-key"
|
||||
serviceName: "search-service.dns"
|
||||
servicePort: 8080
|
||||
```
|
||||
|
||||
### Search Rewrite Configuration
|
||||
|
||||
```yaml
|
||||
searchFrom:
|
||||
- type: google
|
||||
apiKey: "your-google-api-key"
|
||||
cx: "search-engine-id"
|
||||
serviceName: "google-svc.dns"
|
||||
servicePort: 443
|
||||
searchRewrite:
|
||||
llmServiceName: "llm-svc.dns"
|
||||
llmServicePort: 443
|
||||
llmApiKey: "your-llm-api-key"
|
||||
llmUrl: "https://api.example.com/v1/chat/completions"
|
||||
llmModelName: "gpt-3.5-turbo"
|
||||
timeoutMillisecond: 15000
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
1. The prompt template must include `{search_results}` and `{question}` placeholders, optionally use `{cur_date}` to insert current date (format: January 2, 2006)
|
||||
2. The default template includes search results processing instructions and response specifications, you can use the default template unless there are special needs
|
||||
3. Multiple search engines query in parallel, total timeout = maximum timeoutMillisecond value among all search engine configurations + processing time
|
||||
4. Arxiv search doesn't require API key, but you can specify paper category (arxivCategory) to narrow search scope
|
||||
134
plugins/wasm-go/extensions/ai-search/engine/arxiv/arxiv.go
Normal file
134
plugins/wasm-go/extensions/ai-search/engine/arxiv/arxiv.go
Normal file
@@ -0,0 +1,134 @@
|
||||
package arxiv
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
"github.com/antchfx/xmlquery"
|
||||
"github.com/tidwall/gjson"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine"
|
||||
)
|
||||
|
||||
type ArxivSearch struct {
|
||||
optionArgs map[string]string
|
||||
start int
|
||||
count int
|
||||
timeoutMillisecond uint32
|
||||
client wrapper.HttpClient
|
||||
arxivCategory string
|
||||
}
|
||||
|
||||
func NewArxivSearch(config *gjson.Result) (*ArxivSearch, error) {
|
||||
engine := &ArxivSearch{}
|
||||
serviceName := config.Get("serviceName").String()
|
||||
if serviceName == "" {
|
||||
return nil, errors.New("serviceName not found")
|
||||
}
|
||||
servicePort := config.Get("servicePort").Int()
|
||||
if servicePort == 0 {
|
||||
return nil, errors.New("servicePort not found")
|
||||
}
|
||||
engine.client = wrapper.NewClusterClient(wrapper.FQDNCluster{
|
||||
FQDN: serviceName,
|
||||
Port: servicePort,
|
||||
})
|
||||
engine.start = int(config.Get("start").Uint())
|
||||
engine.count = int(config.Get("count").Uint())
|
||||
if engine.count == 0 {
|
||||
engine.count = 10
|
||||
}
|
||||
engine.timeoutMillisecond = uint32(config.Get("timeoutMillisecond").Uint())
|
||||
if engine.timeoutMillisecond == 0 {
|
||||
engine.timeoutMillisecond = 5000
|
||||
}
|
||||
engine.optionArgs = map[string]string{}
|
||||
for key, value := range config.Get("optionArgs").Map() {
|
||||
valStr := value.String()
|
||||
if valStr != "" {
|
||||
engine.optionArgs[key] = value.String()
|
||||
}
|
||||
}
|
||||
engine.arxivCategory = config.Get("arxivCategory").String()
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
func (a ArxivSearch) NeedExectue(ctx engine.SearchContext) bool {
|
||||
return ctx.EngineType == "arxiv"
|
||||
}
|
||||
|
||||
func (a ArxivSearch) Client() wrapper.HttpClient {
|
||||
return a.client
|
||||
}
|
||||
|
||||
func (a ArxivSearch) CallArgs(ctx engine.SearchContext) engine.CallArgs {
|
||||
var searchQueryItems []string
|
||||
for _, q := range ctx.Querys {
|
||||
searchQueryItems = append(searchQueryItems, fmt.Sprintf("all:%s", url.QueryEscape(q)))
|
||||
}
|
||||
searchQuery := strings.Join(searchQueryItems, "+AND+")
|
||||
category := ctx.ArxivCategory
|
||||
if category == "" {
|
||||
category = a.arxivCategory
|
||||
}
|
||||
if category != "" {
|
||||
searchQuery = fmt.Sprintf("%s+AND+cat:%s", searchQuery, category)
|
||||
}
|
||||
queryUrl := fmt.Sprintf("https://export.arxiv.org/api/query?search_query=%s&max_results=%d&start=%d",
|
||||
searchQuery, a.count, a.start)
|
||||
var extraArgs []string
|
||||
for key, value := range a.optionArgs {
|
||||
extraArgs = append(extraArgs, fmt.Sprintf("%s=%s", key, url.QueryEscape(value)))
|
||||
}
|
||||
if len(extraArgs) > 0 {
|
||||
queryUrl = fmt.Sprintf("%s&%s", queryUrl, strings.Join(extraArgs, "&"))
|
||||
}
|
||||
return engine.CallArgs{
|
||||
Method: http.MethodGet,
|
||||
Url: queryUrl,
|
||||
Headers: [][2]string{{"Accept", "application/atom+xml"}},
|
||||
TimeoutMillisecond: a.timeoutMillisecond,
|
||||
}
|
||||
}
|
||||
|
||||
func (a ArxivSearch) ParseResult(ctx engine.SearchContext, response []byte) []engine.SearchResult {
|
||||
var results []engine.SearchResult
|
||||
doc, err := xmlquery.Parse(bytes.NewReader(response))
|
||||
if err != nil {
|
||||
return results
|
||||
}
|
||||
|
||||
entries := xmlquery.Find(doc, "//entry")
|
||||
for _, entry := range entries {
|
||||
title := entry.SelectElement("title").InnerText()
|
||||
link := ""
|
||||
for _, l := range entry.SelectElements("link") {
|
||||
if l.SelectAttr("rel") == "alternate" && l.SelectAttr("type") == "text/html" {
|
||||
link = l.SelectAttr("href")
|
||||
break
|
||||
}
|
||||
}
|
||||
summary := entry.SelectElement("summary").InnerText()
|
||||
publishTime := entry.SelectElement("published").InnerText()
|
||||
authors := entry.SelectElements("author")
|
||||
var authorNames []string
|
||||
for _, author := range authors {
|
||||
authorNames = append(authorNames, author.SelectElement("name").InnerText())
|
||||
}
|
||||
content := fmt.Sprintf("%s\nAuthors: %s\nPublication time: %s", summary, strings.Join(authorNames, ", "), publishTime)
|
||||
result := engine.SearchResult{
|
||||
Title: title,
|
||||
Link: link,
|
||||
Content: content,
|
||||
}
|
||||
if result.Valid() {
|
||||
results = append(results, result)
|
||||
}
|
||||
}
|
||||
return results
|
||||
}
|
||||
128
plugins/wasm-go/extensions/ai-search/engine/bing/bing.go
Normal file
128
plugins/wasm-go/extensions/ai-search/engine/bing/bing.go
Normal file
@@ -0,0 +1,128 @@
|
||||
package bing
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
"github.com/tidwall/gjson"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine"
|
||||
)
|
||||
|
||||
type BingSearch struct {
|
||||
optionArgs map[string]string
|
||||
apiKey string
|
||||
start int
|
||||
count int
|
||||
timeoutMillisecond uint32
|
||||
client wrapper.HttpClient
|
||||
}
|
||||
|
||||
func NewBingSearch(config *gjson.Result) (*BingSearch, error) {
|
||||
engine := &BingSearch{}
|
||||
engine.apiKey = config.Get("apiKey").String()
|
||||
if engine.apiKey == "" {
|
||||
return nil, errors.New("apiKey not found")
|
||||
}
|
||||
serviceName := config.Get("serviceName").String()
|
||||
if serviceName == "" {
|
||||
return nil, errors.New("serviceName not found")
|
||||
}
|
||||
servicePort := config.Get("servicePort").Int()
|
||||
if servicePort == 0 {
|
||||
return nil, errors.New("servicePort not found")
|
||||
}
|
||||
engine.client = wrapper.NewClusterClient(wrapper.FQDNCluster{
|
||||
FQDN: serviceName,
|
||||
Port: servicePort,
|
||||
})
|
||||
engine.start = int(config.Get("start").Uint())
|
||||
engine.count = int(config.Get("count").Uint())
|
||||
if engine.count == 0 {
|
||||
engine.count = 10
|
||||
}
|
||||
engine.timeoutMillisecond = uint32(config.Get("timeoutMillisecond").Uint())
|
||||
if engine.timeoutMillisecond == 0 {
|
||||
engine.timeoutMillisecond = 5000
|
||||
}
|
||||
engine.optionArgs = map[string]string{}
|
||||
for key, value := range config.Get("optionArgs").Map() {
|
||||
valStr := value.String()
|
||||
if valStr != "" {
|
||||
engine.optionArgs[key] = value.String()
|
||||
}
|
||||
}
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
func (b BingSearch) NeedExectue(ctx engine.SearchContext) bool {
|
||||
return ctx.EngineType == "internet"
|
||||
}
|
||||
|
||||
func (b BingSearch) Client() wrapper.HttpClient {
|
||||
return b.client
|
||||
}
|
||||
|
||||
func (b BingSearch) CallArgs(ctx engine.SearchContext) engine.CallArgs {
|
||||
queryUrl := fmt.Sprintf("https://api.bing.microsoft.com/v7.0/search?q=%s&count=%d&offset=%d",
|
||||
url.QueryEscape(strings.Join(ctx.Querys, " ")), b.count, b.start)
|
||||
var extraArgs []string
|
||||
for key, value := range b.optionArgs {
|
||||
extraArgs = append(extraArgs, fmt.Sprintf("%s=%s", key, url.QueryEscape(value)))
|
||||
}
|
||||
if ctx.Language != "" {
|
||||
extraArgs = append(extraArgs, fmt.Sprintf("mkt=%s", ctx.Language))
|
||||
}
|
||||
if len(extraArgs) > 0 {
|
||||
queryUrl = fmt.Sprintf("%s&%s", queryUrl, strings.Join(extraArgs, "&"))
|
||||
}
|
||||
return engine.CallArgs{
|
||||
Method: http.MethodGet,
|
||||
Url: queryUrl,
|
||||
Headers: [][2]string{{"Ocp-Apim-Subscription-Key", b.apiKey}},
|
||||
TimeoutMillisecond: b.timeoutMillisecond,
|
||||
}
|
||||
}
|
||||
|
||||
func (b BingSearch) ParseResult(ctx engine.SearchContext, response []byte) []engine.SearchResult {
|
||||
jsonObj := gjson.ParseBytes(response)
|
||||
var results []engine.SearchResult
|
||||
webPages := jsonObj.Get("webPages.value")
|
||||
for _, page := range webPages.Array() {
|
||||
result := engine.SearchResult{
|
||||
Title: page.Get("name").String(),
|
||||
Link: page.Get("url").String(),
|
||||
Content: page.Get("snippet").String(),
|
||||
}
|
||||
if result.Valid() {
|
||||
results = append(results, result)
|
||||
}
|
||||
deepLinks := page.Get("deepLinks")
|
||||
for _, inner := range deepLinks.Array() {
|
||||
innerResult := engine.SearchResult{
|
||||
Title: inner.Get("name").String(),
|
||||
Link: inner.Get("url").String(),
|
||||
Content: inner.Get("snippet").String(),
|
||||
}
|
||||
if innerResult.Valid() {
|
||||
results = append(results, innerResult)
|
||||
}
|
||||
}
|
||||
}
|
||||
news := jsonObj.Get("news.value")
|
||||
for _, article := range news.Array() {
|
||||
result := engine.SearchResult{
|
||||
Title: article.Get("name").String(),
|
||||
Link: article.Get("url").String(),
|
||||
Content: article.Get("description").String(),
|
||||
}
|
||||
if result.Valid() {
|
||||
results = append(results, result)
|
||||
}
|
||||
}
|
||||
return results
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
package elasticsearch
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
"github.com/tidwall/gjson"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine"
|
||||
)
|
||||
|
||||
type ElasticsearchSearch struct {
|
||||
client wrapper.HttpClient
|
||||
index string
|
||||
contentField string
|
||||
linkField string
|
||||
titleField string
|
||||
start int
|
||||
count int
|
||||
timeoutMillisecond uint32
|
||||
}
|
||||
|
||||
func NewElasticsearchSearch(config *gjson.Result) (*ElasticsearchSearch, error) {
|
||||
engine := &ElasticsearchSearch{}
|
||||
serviceName := config.Get("serviceName").String()
|
||||
if serviceName == "" {
|
||||
return nil, errors.New("serviceName not found")
|
||||
}
|
||||
servicePort := config.Get("servicePort").Int()
|
||||
if servicePort == 0 {
|
||||
return nil, errors.New("servicePort not found")
|
||||
}
|
||||
engine.client = wrapper.NewClusterClient(wrapper.FQDNCluster{
|
||||
FQDN: serviceName,
|
||||
Port: servicePort,
|
||||
})
|
||||
engine.index = config.Get("index").String()
|
||||
if engine.index == "" {
|
||||
return nil, errors.New("index not found")
|
||||
}
|
||||
engine.contentField = config.Get("contentField").String()
|
||||
if engine.contentField == "" {
|
||||
return nil, errors.New("contentField not found")
|
||||
}
|
||||
engine.linkField = config.Get("linkField").String()
|
||||
if engine.linkField == "" {
|
||||
return nil, errors.New("linkField not found")
|
||||
}
|
||||
engine.titleField = config.Get("titleField").String()
|
||||
if engine.titleField == "" {
|
||||
return nil, errors.New("titleField not found")
|
||||
}
|
||||
engine.timeoutMillisecond = uint32(config.Get("timeoutMillisecond").Uint())
|
||||
if engine.timeoutMillisecond == 0 {
|
||||
engine.timeoutMillisecond = 5000
|
||||
}
|
||||
engine.start = int(config.Get("start").Uint())
|
||||
engine.count = int(config.Get("count").Uint())
|
||||
if engine.count == 0 {
|
||||
engine.count = 10
|
||||
}
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
func (e ElasticsearchSearch) NeedExectue(ctx engine.SearchContext) bool {
|
||||
return ctx.EngineType == "private"
|
||||
}
|
||||
|
||||
func (e ElasticsearchSearch) Client() wrapper.HttpClient {
|
||||
return e.client
|
||||
}
|
||||
|
||||
func (e ElasticsearchSearch) CallArgs(ctx engine.SearchContext) engine.CallArgs {
|
||||
searchBody := fmt.Sprintf(`{
|
||||
"query": {
|
||||
"match": {
|
||||
"%s": {
|
||||
"query": "%s",
|
||||
"operator": "AND"
|
||||
}
|
||||
}
|
||||
}
|
||||
}`, e.contentField, strings.Join(ctx.Querys, " "))
|
||||
|
||||
return engine.CallArgs{
|
||||
Method: http.MethodPost,
|
||||
Url: fmt.Sprintf("/%s/_search?from=%d&size=%d", e.index, e.start, e.count),
|
||||
Headers: [][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
},
|
||||
Body: []byte(searchBody),
|
||||
TimeoutMillisecond: e.timeoutMillisecond,
|
||||
}
|
||||
}
|
||||
|
||||
func (e ElasticsearchSearch) ParseResult(ctx engine.SearchContext, response []byte) []engine.SearchResult {
|
||||
jsonObj := gjson.ParseBytes(response)
|
||||
var results []engine.SearchResult
|
||||
for _, hit := range jsonObj.Get("hits.hits").Array() {
|
||||
source := hit.Get("_source")
|
||||
result := engine.SearchResult{
|
||||
Title: source.Get(e.titleField).String(),
|
||||
Link: source.Get(e.linkField).String(),
|
||||
Content: source.Get(e.contentField).String(),
|
||||
}
|
||||
if result.Valid() {
|
||||
results = append(results, result)
|
||||
}
|
||||
}
|
||||
return results
|
||||
}
|
||||
120
plugins/wasm-go/extensions/ai-search/engine/google/google.go
Normal file
120
plugins/wasm-go/extensions/ai-search/engine/google/google.go
Normal file
@@ -0,0 +1,120 @@
|
||||
package google
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"strings"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
"github.com/tidwall/gjson"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine"
|
||||
)
|
||||
|
||||
type GoogleSearch struct {
|
||||
optionArgs map[string]string
|
||||
apiKey string
|
||||
cx string
|
||||
start int
|
||||
count int
|
||||
timeoutMillisecond uint32
|
||||
client wrapper.HttpClient
|
||||
}
|
||||
|
||||
func NewGoogleSearch(config *gjson.Result) (*GoogleSearch, error) {
|
||||
engine := &GoogleSearch{}
|
||||
engine.apiKey = config.Get("apiKey").String()
|
||||
if engine.apiKey == "" {
|
||||
return nil, errors.New("apiKey not found")
|
||||
}
|
||||
engine.cx = config.Get("cx").String()
|
||||
if engine.cx == "" {
|
||||
return nil, errors.New("cx not found")
|
||||
}
|
||||
serviceName := config.Get("serviceName").String()
|
||||
if serviceName == "" {
|
||||
return nil, errors.New("serviceName not found")
|
||||
}
|
||||
servicePort := config.Get("servicePort").Int()
|
||||
if servicePort == 0 {
|
||||
return nil, errors.New("servicePort not found")
|
||||
}
|
||||
engine.client = wrapper.NewClusterClient(wrapper.FQDNCluster{
|
||||
FQDN: serviceName,
|
||||
Port: servicePort,
|
||||
})
|
||||
engine.start = int(config.Get("start").Uint())
|
||||
engine.count = int(config.Get("count").Uint())
|
||||
if engine.count == 0 {
|
||||
engine.count = 10
|
||||
}
|
||||
if engine.count > 10 || engine.start+engine.count > 100 {
|
||||
return nil, errors.New("count must be less than 10, and start + count must be less than or equal to 100.")
|
||||
}
|
||||
engine.timeoutMillisecond = uint32(config.Get("timeoutMillisecond").Uint())
|
||||
if engine.timeoutMillisecond == 0 {
|
||||
engine.timeoutMillisecond = 5000
|
||||
}
|
||||
engine.optionArgs = map[string]string{}
|
||||
for key, value := range config.Get("optionArgs").Map() {
|
||||
valStr := value.String()
|
||||
if valStr != "" {
|
||||
engine.optionArgs[key] = value.String()
|
||||
}
|
||||
}
|
||||
return engine, nil
|
||||
}
|
||||
|
||||
func (g GoogleSearch) NeedExectue(ctx engine.SearchContext) bool {
|
||||
return ctx.EngineType == "internet"
|
||||
}
|
||||
|
||||
func (g GoogleSearch) Client() wrapper.HttpClient {
|
||||
return g.client
|
||||
}
|
||||
|
||||
func (g GoogleSearch) CallArgs(ctx engine.SearchContext) engine.CallArgs {
|
||||
queryUrl := fmt.Sprintf("https://customsearch.googleapis.com/customsearch/v1?cx=%s&q=%s&num=%d&key=%s&start=%d",
|
||||
g.cx, url.QueryEscape(strings.Join(ctx.Querys, " ")), g.count, g.apiKey, g.start+1)
|
||||
var extraArgs []string
|
||||
for key, value := range g.optionArgs {
|
||||
extraArgs = append(extraArgs, fmt.Sprintf("%s=%s", key, url.QueryEscape(value)))
|
||||
}
|
||||
if ctx.Language != "" {
|
||||
extraArgs = append(extraArgs, fmt.Sprintf("lr=lang_%s", ctx.Language))
|
||||
}
|
||||
if len(extraArgs) > 0 {
|
||||
queryUrl = fmt.Sprintf("%s&%s", queryUrl, strings.Join(extraArgs, "&"))
|
||||
}
|
||||
return engine.CallArgs{
|
||||
Method: http.MethodGet,
|
||||
Url: queryUrl,
|
||||
Headers: [][2]string{
|
||||
{"Accept", "application/json"},
|
||||
},
|
||||
TimeoutMillisecond: g.timeoutMillisecond,
|
||||
}
|
||||
}
|
||||
|
||||
func (g GoogleSearch) ParseResult(ctx engine.SearchContext, response []byte) []engine.SearchResult {
|
||||
jsonObj := gjson.ParseBytes(response)
|
||||
var results []engine.SearchResult
|
||||
for _, item := range jsonObj.Get("items").Array() {
|
||||
content := item.Get("snippet").String()
|
||||
metaDescription := item.Get("pagemap.metatags.0.og:description").String()
|
||||
if metaDescription != "" {
|
||||
content = fmt.Sprintf("%s\n...\n%s", content, metaDescription)
|
||||
}
|
||||
result := engine.SearchResult{
|
||||
Title: item.Get("title").String(),
|
||||
Link: item.Get("link").String(),
|
||||
Content: content,
|
||||
}
|
||||
if result.Valid() {
|
||||
results = append(results, result)
|
||||
}
|
||||
}
|
||||
return results
|
||||
}
|
||||
37
plugins/wasm-go/extensions/ai-search/engine/types.go
Normal file
37
plugins/wasm-go/extensions/ai-search/engine/types.go
Normal file
@@ -0,0 +1,37 @@
|
||||
package engine
|
||||
|
||||
import (
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
)
|
||||
|
||||
type SearchResult struct {
|
||||
Title string
|
||||
Link string
|
||||
Content string
|
||||
}
|
||||
|
||||
func (result SearchResult) Valid() bool {
|
||||
return result.Title != "" && result.Link != "" && result.Content != ""
|
||||
}
|
||||
|
||||
type SearchContext struct {
|
||||
EngineType string
|
||||
Querys []string
|
||||
Language string
|
||||
ArxivCategory string
|
||||
}
|
||||
|
||||
type CallArgs struct {
|
||||
Method string
|
||||
Url string
|
||||
Headers [][2]string
|
||||
Body []byte
|
||||
TimeoutMillisecond uint32
|
||||
}
|
||||
|
||||
type SearchEngine interface {
|
||||
NeedExectue(ctx SearchContext) bool
|
||||
Client() wrapper.HttpClient
|
||||
CallArgs(ctx SearchContext) CallArgs
|
||||
ParseResult(ctx SearchContext, response []byte) []SearchResult
|
||||
}
|
||||
26
plugins/wasm-go/extensions/ai-search/go.mod
Normal file
26
plugins/wasm-go/extensions/ai-search/go.mod
Normal file
@@ -0,0 +1,26 @@
|
||||
module github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search
|
||||
|
||||
go 1.18
|
||||
|
||||
replace github.com/alibaba/higress/plugins/wasm-go => ../..
|
||||
|
||||
require (
|
||||
github.com/alibaba/higress/plugins/wasm-go v0.0.0
|
||||
github.com/antchfx/xmlquery v1.4.4
|
||||
github.com/higress-group/proxy-wasm-go-sdk v1.0.0
|
||||
github.com/tidwall/gjson v1.18.0
|
||||
github.com/tidwall/sjson v1.2.5
|
||||
)
|
||||
|
||||
require (
|
||||
github.com/antchfx/xpath v1.3.3 // indirect
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
|
||||
github.com/google/uuid v1.3.0 // indirect
|
||||
github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 // indirect
|
||||
github.com/magefile/mage v1.14.0 // indirect
|
||||
github.com/tidwall/match v1.1.1 // indirect
|
||||
github.com/tidwall/pretty v1.2.0 // indirect
|
||||
github.com/tidwall/resp v0.1.1 // indirect
|
||||
golang.org/x/net v0.33.0 // indirect
|
||||
golang.org/x/text v0.21.0 // indirect
|
||||
)
|
||||
96
plugins/wasm-go/extensions/ai-search/go.sum
Normal file
96
plugins/wasm-go/extensions/ai-search/go.sum
Normal file
@@ -0,0 +1,96 @@
|
||||
github.com/antchfx/xmlquery v1.4.4 h1:mxMEkdYP3pjKSftxss4nUHfjBhnMk4imGoR96FRY2dg=
|
||||
github.com/antchfx/xmlquery v1.4.4/go.mod h1:AEPEEPYE9GnA2mj5Ur2L5Q5/2PycJ0N9Fusrx9b12fc=
|
||||
github.com/antchfx/xpath v1.3.3 h1:tmuPQa1Uye0Ym1Zn65vxPgfltWb/Lxu2jeqIGteJSRs=
|
||||
github.com/antchfx/xpath v1.3.3/go.mod h1:i54GszH55fYfBmoZXapTHN8T8tkcHfRgLyVwwqzXNcs=
|
||||
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da h1:oI5xCqsCo564l8iNU+DwB5epxmsaqB+rhGL0m5jtYqE=
|
||||
github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
|
||||
github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
|
||||
github.com/google/uuid v1.3.0 h1:t6JiXgmwXMjEs8VusXIJk2BXHsn+wx8BZdTaoZ5fu7I=
|
||||
github.com/google/uuid v1.3.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
|
||||
github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520 h1:IHDghbGQ2DTIXHBHxWfqCYQW1fKjyJ/I7W1pMyUDeEA=
|
||||
github.com/higress-group/nottinygc v0.0.0-20231101025119-e93c4c2f8520/go.mod h1:Nz8ORLaFiLWotg6GeKlJMhv8cci8mM43uEnLA5t8iew=
|
||||
github.com/higress-group/proxy-wasm-go-sdk v1.0.0 h1:BZRNf4R7jr9hwRivg/E29nkVaKEak5MWjBDhWjuHijU=
|
||||
github.com/higress-group/proxy-wasm-go-sdk v1.0.0/go.mod h1:iiSyFbo+rAtbtGt/bsefv8GU57h9CCLYGJA74/tF5/0=
|
||||
github.com/magefile/mage v1.14.0 h1:6QDX3g6z1YvJ4olPhT1wksUcSa/V0a1B+pJb73fBjyo=
|
||||
github.com/magefile/mage v1.14.0/go.mod h1:z5UZb/iS3GoOSn0JgWuiw7dxlurVYTu+/jHXqQg881A=
|
||||
github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
|
||||
github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
|
||||
github.com/tidwall/gjson v1.14.2/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
|
||||
github.com/tidwall/gjson v1.18.0 h1:FIDeeyB800efLX89e5a8Y0BNH+LOngJyGrIWxG2FKQY=
|
||||
github.com/tidwall/gjson v1.18.0/go.mod h1:/wbyibRr2FHMks5tjHJ5F8dMZh3AcwJEMf5vlfC0lxk=
|
||||
github.com/tidwall/match v1.1.1 h1:+Ho715JplO36QYgwN9PGYNhgZvoUSc9X2c80KVTi+GA=
|
||||
github.com/tidwall/match v1.1.1/go.mod h1:eRSPERbgtNPcGhD8UCthc6PmLEQXEWd3PRB5JTxsfmM=
|
||||
github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs=
|
||||
github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU=
|
||||
github.com/tidwall/resp v0.1.1 h1:Ly20wkhqKTmDUPlyM1S7pWo5kk0tDu8OoC/vFArXmwE=
|
||||
github.com/tidwall/resp v0.1.1/go.mod h1:3/FrruOBAxPTPtundW0VXgmsQ4ZBA0Aw714lVYgwFa0=
|
||||
github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY=
|
||||
github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28=
|
||||
github.com/yuin/goldmark v1.4.13/go.mod h1:6yULJ656Px+3vBD8DxQVa3kxgyrAnzto9xy5taEt/CY=
|
||||
golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
|
||||
golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
|
||||
golang.org/x/crypto v0.13.0/go.mod h1:y6Z2r+Rw4iayiXXAIxJIDAJ1zMW4yaTpebo8fPOliYc=
|
||||
golang.org/x/crypto v0.19.0/go.mod h1:Iy9bg/ha4yyC70EfRS8jz+B6ybOBKMaSxLj6P6oBDfU=
|
||||
golang.org/x/crypto v0.23.0/go.mod h1:CKFgDieR+mRhux2Lsu27y0fO304Db0wZe70UKqHu0v8=
|
||||
golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk=
|
||||
golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
|
||||
golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
|
||||
golang.org/x/mod v0.15.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/mod v0.17.0/go.mod h1:hTbmBsO62+eylJbnUtE2MGJUyE7QWk4xUqPFrRgJ+7c=
|
||||
golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
|
||||
golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v0D8zg8gWTRqZa9RBIspLL5mdg=
|
||||
golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug+ECip1KBveYUHfp+8e9klMJ9c=
|
||||
golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
|
||||
golang.org/x/net v0.10.0/go.mod h1:0qNGK6F8kojg2nk9dLZ2mShWaEBan6FAoqfSigmmuDg=
|
||||
golang.org/x/net v0.15.0/go.mod h1:idbUs1IY1+zTqbi8yxTbhexhEEk5ur9LInksu6HrEpk=
|
||||
golang.org/x/net v0.21.0/go.mod h1:bIjVDfnllIU7BJ2DNgfnXvpSvtn8VRwhlsaeUTyUS44=
|
||||
golang.org/x/net v0.25.0/go.mod h1:JkAGAh7GEvH74S6FOH42FLoXpXbE/aqXSrIQjXgsiwM=
|
||||
golang.org/x/net v0.33.0 h1:74SYHlV8BIgHIFC/LrYkOGIwL19eTYXQ5wc6TBuO36I=
|
||||
golang.org/x/net v0.33.0/go.mod h1:HXLR5J+9DxmrqMwG9qjGCxZ+zKXxBru04zlTvWlWuN4=
|
||||
golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
|
||||
golang.org/x/sync v0.3.0/go.mod h1:FU7BRWz2tNW+3quACPkgCx/L+uEAv1htQ0V83Z9Rj+Y=
|
||||
golang.org/x/sync v0.6.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.7.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk=
|
||||
golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
|
||||
golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
|
||||
golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220520151302-bc2c85ada10a/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.12.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
|
||||
golang.org/x/sys v0.17.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.20.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
|
||||
golang.org/x/telemetry v0.0.0-20240228155512-f48c80bd79b2/go.mod h1:TeRTkGYfJXctD9OcfyVLyj2J3IxLnKwHJR8f4D8a3YE=
|
||||
golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
|
||||
golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
|
||||
golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
|
||||
golang.org/x/term v0.8.0/go.mod h1:xPskH00ivmX89bAKVGSKKtLOWNx2+17Eiy94tnKShWo=
|
||||
golang.org/x/term v0.12.0/go.mod h1:owVbMEjm3cBLCHdkQu9b1opXd4ETQWc3BhuQGKgXgvU=
|
||||
golang.org/x/term v0.17.0/go.mod h1:lLRBjIVuehSbZlaOtGMbcMncT+aqLLLmKrsjNrUguwk=
|
||||
golang.org/x/term v0.20.0/go.mod h1:8UkIAJTvZgivsXaD6/pH6U9ecQzZ45awqEOzuCvwpFY=
|
||||
golang.org/x/term v0.27.0/go.mod h1:iMsnZpn0cago0GOrHO2+Y7u7JPn5AylBrcoWkElMTSM=
|
||||
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
|
||||
golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
|
||||
golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
|
||||
golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8=
|
||||
golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8=
|
||||
golang.org/x/text v0.13.0/go.mod h1:TvPlkZtksWOMsz7fbANvkp4WM8x/WCo/om8BMLbz+aE=
|
||||
golang.org/x/text v0.14.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.15.0/go.mod h1:18ZOQIKpY8NJVqYksKHtTdi31H5itFRjB5/qKTNYzSU=
|
||||
golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
|
||||
golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
|
||||
golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
|
||||
golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=
|
||||
golang.org/x/tools v0.1.12/go.mod h1:hNGJHUnrk76NpqgfD5Aqm5Crs+Hm0VOH/i9J2+nxYbc=
|
||||
golang.org/x/tools v0.6.0/go.mod h1:Xwgl3UAJ/d3gWutnCtw505GrjyAbvKui8lOU390QaIU=
|
||||
golang.org/x/tools v0.13.0/go.mod h1:HvlwmtVNQAhOuCjW7xxvovg8wbNq7LwfXh/k7wXUl58=
|
||||
golang.org/x/tools v0.21.1-0.20240508182429-e35e4ccd0d2d/go.mod h1:aiJjzUbINMkxbQROHiO6hDPo2LHcIPhhQsa9DLh0yGk=
|
||||
golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
|
||||
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
|
||||
559
plugins/wasm-go/extensions/ai-search/main.go
Normal file
559
plugins/wasm-go/extensions/ai-search/main.go
Normal file
@@ -0,0 +1,559 @@
|
||||
// Copyright (c) 2022 Alibaba Group Holding Ltd.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package main
|
||||
|
||||
import (
|
||||
_ "embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm"
|
||||
"github.com/higress-group/proxy-wasm-go-sdk/proxywasm/types"
|
||||
"github.com/tidwall/gjson"
|
||||
"github.com/tidwall/sjson"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/pkg/wrapper"
|
||||
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine"
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/arxiv"
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/bing"
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/elasticsearch"
|
||||
"github.com/alibaba/higress/plugins/wasm-go/extensions/ai-search/engine/google"
|
||||
)
|
||||
|
||||
type SearchRewrite struct {
|
||||
client wrapper.HttpClient
|
||||
url string
|
||||
apiKey string
|
||||
modelName string
|
||||
timeoutMillisecond uint32
|
||||
prompt string
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
engine []engine.SearchEngine
|
||||
promptTemplate string
|
||||
referenceFormat string
|
||||
defaultLanguage string
|
||||
needReference bool
|
||||
searchRewrite *SearchRewrite
|
||||
}
|
||||
|
||||
const (
|
||||
DEFAULT_MAX_BODY_BYTES uint32 = 100 * 1024 * 1024
|
||||
)
|
||||
|
||||
//go:embed prompts/full.md
|
||||
var fullSearchPrompts string
|
||||
|
||||
//go:embed prompts/arxiv.md
|
||||
var arxivSearchPrompts string
|
||||
|
||||
//go:embed prompts/internet.md
|
||||
var internetSearchPrompts string
|
||||
|
||||
//go:embed prompts/private.md
|
||||
var privateSearchPrompts string
|
||||
|
||||
func main() {
|
||||
wrapper.SetCtx(
|
||||
"ai-search",
|
||||
wrapper.ParseConfigBy(parseConfig),
|
||||
wrapper.ProcessRequestHeadersBy(onHttpRequestHeaders),
|
||||
wrapper.ProcessRequestBodyBy(onHttpRequestBody),
|
||||
wrapper.ProcessResponseHeadersBy(onHttpResponseHeaders),
|
||||
wrapper.ProcessStreamingResponseBodyBy(onStreamingResponseBody),
|
||||
wrapper.ProcessResponseBodyBy(onHttpResponseBody),
|
||||
)
|
||||
}
|
||||
|
||||
func parseConfig(json gjson.Result, config *Config, log wrapper.Log) error {
|
||||
config.needReference = json.Get("needReference").Bool()
|
||||
if config.needReference {
|
||||
config.referenceFormat = json.Get("referenceFormat").String()
|
||||
if config.referenceFormat == "" {
|
||||
config.referenceFormat = "**References:**\n%s"
|
||||
} else if !strings.Contains(config.referenceFormat, "%s") {
|
||||
return fmt.Errorf("invalid referenceFormat:%s", config.referenceFormat)
|
||||
}
|
||||
}
|
||||
config.defaultLanguage = json.Get("defaultLang").String()
|
||||
config.promptTemplate = json.Get("promptTemplate").String()
|
||||
if config.promptTemplate == "" {
|
||||
if config.needReference {
|
||||
config.promptTemplate = `# 以下内容是基于用户发送的消息的搜索结果:
|
||||
{search_results}
|
||||
在我给你的搜索结果中，每个结果都是[webpage X begin]...[webpage X end]格式的，X代表每篇文章的数字索引。请在适当的情况下在句子末尾引用上下文。请按照引用编号[X]的格式在答案中对应部分引用上下文。如果一句话源自多个上下文，请列出所有相关的引用编号，例如[3][5]，切记不要将引用集中在最后返回引用编号，而是在答案对应部分列出。
|
||||
在回答时，请注意以下几点：
|
||||
- 今天是北京时间：{cur_date}。
|
||||
- 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
|
||||
- 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。
|
||||
- 对于创作类的问题（如写论文），请务必在正文的段落中引用对应的参考编号，例如[3][5]，不能只在文章末尾引用。你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。
|
||||
- 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
|
||||
- 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
|
||||
- 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
|
||||
- 你的回答应该综合多个相关网页来回答，不能重复引用一个网页。
|
||||
- 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
|
||||
|
||||
# 用户消息为：
|
||||
{question}`
|
||||
} else {
|
||||
config.promptTemplate = `# 以下内容是基于用户发送的消息的搜索结果:
|
||||
{search_results}
|
||||
在我给你的搜索结果中，每个结果都是[webpage begin]...[webpage end]格式的。
|
||||
在回答时，请注意以下几点：
|
||||
- 今天是北京时间：{cur_date}。
|
||||
- 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
|
||||
- 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内。如非必要，不要主动告诉用户搜索结果未提供的内容。
|
||||
- 对于创作类的问题（如写论文），你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。
|
||||
- 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
|
||||
- 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
|
||||
- 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
|
||||
- 你的回答应该综合多个相关网页来回答，但回答中不要给出网页的引用来源。
|
||||
- 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
|
||||
|
||||
# 用户消息为：
|
||||
{question}`
|
||||
}
|
||||
}
|
||||
if !strings.Contains(config.promptTemplate, "{search_results}") ||
|
||||
!strings.Contains(config.promptTemplate, "{question}") {
|
||||
return fmt.Errorf("invalid promptTemplate, must contains {search_results} and {question}:%s", config.promptTemplate)
|
||||
}
|
||||
var internetExists, privateExists, arxivExists bool
|
||||
for _, e := range json.Get("searchFrom").Array() {
|
||||
switch e.Get("type").String() {
|
||||
case "being":
|
||||
searchEngine, err := bing.NewBingSearch(&e)
|
||||
if err != nil {
|
||||
return fmt.Errorf("being search engine init failed:%s", err)
|
||||
}
|
||||
config.engine = append(config.engine, searchEngine)
|
||||
internetExists = true
|
||||
case "google":
|
||||
searchEngine, err := google.NewGoogleSearch(&e)
|
||||
if err != nil {
|
||||
return fmt.Errorf("google search engine init failed:%s", err)
|
||||
}
|
||||
config.engine = append(config.engine, searchEngine)
|
||||
internetExists = true
|
||||
case "arxiv":
|
||||
searchEngine, err := arxiv.NewArxivSearch(&e)
|
||||
if err != nil {
|
||||
return fmt.Errorf("arxiv search engine init failed:%s", err)
|
||||
}
|
||||
config.engine = append(config.engine, searchEngine)
|
||||
arxivExists = true
|
||||
case "elasticsearch":
|
||||
searchEngine, err := elasticsearch.NewElasticsearchSearch(&e)
|
||||
if err != nil {
|
||||
return fmt.Errorf("elasticsearch search engine init failed:%s", err)
|
||||
}
|
||||
config.engine = append(config.engine, searchEngine)
|
||||
privateExists = true
|
||||
default:
|
||||
return fmt.Errorf("unkown search engine:%s", e.Get("type").String())
|
||||
}
|
||||
}
|
||||
searchRewriteJson := json.Get("searchRewrite")
|
||||
if searchRewriteJson.Exists() {
|
||||
searchRewrite := &SearchRewrite{}
|
||||
llmServiceName := searchRewriteJson.Get("llmServiceName").String()
|
||||
if llmServiceName == "" {
|
||||
return errors.New("llm_service_name not found")
|
||||
}
|
||||
llmServicePort := searchRewriteJson.Get("llmServicePort").Int()
|
||||
if llmServicePort == 0 {
|
||||
return errors.New("llmServicePort not found")
|
||||
}
|
||||
searchRewrite.client = wrapper.NewClusterClient(wrapper.FQDNCluster{
|
||||
FQDN: llmServiceName,
|
||||
Port: llmServicePort,
|
||||
})
|
||||
llmApiKey := searchRewriteJson.Get("llmApiKey").String()
|
||||
if llmApiKey == "" {
|
||||
return errors.New("llmApiKey not found")
|
||||
}
|
||||
searchRewrite.apiKey = llmApiKey
|
||||
llmUrl := searchRewriteJson.Get("llmUrl").String()
|
||||
if llmUrl == "" {
|
||||
return errors.New("llmUrl not found")
|
||||
}
|
||||
searchRewrite.url = llmUrl
|
||||
llmModelName := searchRewriteJson.Get("llmModelName").String()
|
||||
if llmModelName == "" {
|
||||
return errors.New("llmModelName not found")
|
||||
}
|
||||
searchRewrite.modelName = llmModelName
|
||||
llmTimeout := searchRewriteJson.Get("timeoutMillisecond").Uint()
|
||||
if llmTimeout == 0 {
|
||||
llmTimeout = 30000
|
||||
}
|
||||
searchRewrite.timeoutMillisecond = uint32(llmTimeout)
|
||||
// The consideration here is that internet searches are generally available, but arxiv and private sources may not be.
|
||||
if arxivExists {
|
||||
if privateExists {
|
||||
// private + internet + arxiv
|
||||
searchRewrite.prompt = fullSearchPrompts
|
||||
} else {
|
||||
// internet + arxiv
|
||||
searchRewrite.prompt = arxivSearchPrompts
|
||||
}
|
||||
} else if privateExists {
|
||||
// private + internet
|
||||
searchRewrite.prompt = privateSearchPrompts
|
||||
} else if internetExists {
|
||||
// only internet
|
||||
searchRewrite.prompt = internetSearchPrompts
|
||||
}
|
||||
config.searchRewrite = searchRewrite
|
||||
}
|
||||
if len(config.engine) == 0 {
|
||||
return fmt.Errorf("no avaliable search engine found")
|
||||
}
|
||||
log.Debugf("ai search enabled, config: %#v", config)
|
||||
return nil
|
||||
}
|
||||
|
||||
func onHttpRequestHeaders(ctx wrapper.HttpContext, config Config, log wrapper.Log) types.Action {
|
||||
contentType, _ := proxywasm.GetHttpRequestHeader("content-type")
|
||||
// The request does not have a body.
|
||||
if contentType == "" {
|
||||
return types.ActionContinue
|
||||
}
|
||||
if !strings.Contains(contentType, "application/json") {
|
||||
log.Warnf("content is not json, can't process: %s", contentType)
|
||||
ctx.DontReadRequestBody()
|
||||
return types.ActionContinue
|
||||
}
|
||||
ctx.SetRequestBodyBufferLimit(DEFAULT_MAX_BODY_BYTES)
|
||||
_ = proxywasm.RemoveHttpRequestHeader("Accept-Encoding")
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
func onHttpRequestBody(ctx wrapper.HttpContext, config Config, body []byte, log wrapper.Log) types.Action {
|
||||
var queryIndex int
|
||||
var query string
|
||||
messages := gjson.GetBytes(body, "messages").Array()
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
if messages[i].Get("role").String() == "user" {
|
||||
queryIndex = i
|
||||
query = messages[i].Get("content").String()
|
||||
break
|
||||
}
|
||||
}
|
||||
if query == "" {
|
||||
log.Errorf("not found user query in body:%s", body)
|
||||
return types.ActionContinue
|
||||
}
|
||||
searchRewrite := config.searchRewrite
|
||||
if searchRewrite != nil {
|
||||
startTime := time.Now()
|
||||
rewritePrompt := strings.Replace(searchRewrite.prompt, "{question}", query, 1)
|
||||
rewriteBody, _ := sjson.SetBytes([]byte(fmt.Sprintf(
|
||||
`{"stream":false,"max_tokens":100,"model":"%s","messages":[{"role":"user","content":""}]}`,
|
||||
searchRewrite.modelName)), "messages.0.content", rewritePrompt)
|
||||
err := searchRewrite.client.Post(searchRewrite.url,
|
||||
[][2]string{
|
||||
{"Content-Type", "application/json"},
|
||||
{"Authorization", fmt.Sprintf("Bearer %s", searchRewrite.apiKey)},
|
||||
}, rewriteBody,
|
||||
func(statusCode int, responseHeaders http.Header, responseBody []byte) {
|
||||
if statusCode != http.StatusOK {
|
||||
log.Errorf("search rewrite failed, status: %d", statusCode)
|
||||
// After a rewrite failure, no further search is performed, thus quickly identifying the failure.
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
|
||||
content := gjson.GetBytes(responseBody, "choices.0.message.content").String()
|
||||
log.Infof("LLM rewritten query response: %s (took %v), original search query:%s",
|
||||
strings.ReplaceAll(content, "\n", `\n`), time.Since(startTime), query)
|
||||
if strings.Contains(content, "none") {
|
||||
log.Debugf("no search required")
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
|
||||
// Parse search queries from LLM response
|
||||
var searchContexts []engine.SearchContext
|
||||
for _, line := range strings.Split(content, "\n") {
|
||||
line = strings.TrimSpace(line)
|
||||
if line == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
parts := strings.SplitN(line, ":", 2)
|
||||
if len(parts) != 2 {
|
||||
continue
|
||||
}
|
||||
|
||||
engineType := strings.TrimSpace(parts[0])
|
||||
queryStr := strings.TrimSpace(parts[1])
|
||||
|
||||
var ctx engine.SearchContext
|
||||
ctx.Language = config.defaultLanguage
|
||||
|
||||
switch {
|
||||
case engineType == "internet":
|
||||
ctx.EngineType = engineType
|
||||
ctx.Querys = []string{queryStr}
|
||||
case engineType == "private":
|
||||
ctx.EngineType = engineType
|
||||
ctx.Querys = strings.Split(queryStr, ",")
|
||||
for i := range ctx.Querys {
|
||||
ctx.Querys[i] = strings.TrimSpace(ctx.Querys[i])
|
||||
}
|
||||
default:
|
||||
// Arxiv category
|
||||
ctx.EngineType = "arxiv"
|
||||
ctx.ArxivCategory = engineType
|
||||
ctx.Querys = strings.Split(queryStr, ",")
|
||||
for i := range ctx.Querys {
|
||||
ctx.Querys[i] = strings.TrimSpace(ctx.Querys[i])
|
||||
}
|
||||
}
|
||||
|
||||
if len(ctx.Querys) > 0 {
|
||||
searchContexts = append(searchContexts, ctx)
|
||||
if ctx.ArxivCategory != "" {
|
||||
// Conduct i/nquiries in all areas to increase recall.
|
||||
backupCtx := ctx
|
||||
backupCtx.ArxivCategory = ""
|
||||
searchContexts = append(searchContexts, backupCtx)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if len(searchContexts) == 0 {
|
||||
log.Errorf("no valid search contexts found")
|
||||
proxywasm.ResumeHttpRequest()
|
||||
return
|
||||
}
|
||||
if types.ActionContinue == executeSearch(ctx, config, queryIndex, body, searchContexts, log) {
|
||||
proxywasm.ResumeHttpRequest()
|
||||
}
|
||||
}, searchRewrite.timeoutMillisecond)
|
||||
if err != nil {
|
||||
log.Errorf("search rewrite call llm service failed:%s", err)
|
||||
// After a rewrite failure, no further search is performed, thus quickly identifying the failure.
|
||||
return types.ActionContinue
|
||||
}
|
||||
return types.ActionPause
|
||||
}
|
||||
|
||||
// Execute search without rewrite
|
||||
return executeSearch(ctx, config, queryIndex, body, []engine.SearchContext{{
|
||||
Querys: []string{query},
|
||||
Language: config.defaultLanguage,
|
||||
}}, log)
|
||||
}
|
||||
|
||||
func executeSearch(ctx wrapper.HttpContext, config Config, queryIndex int, body []byte, searchContexts []engine.SearchContext, log wrapper.Log) types.Action {
|
||||
searchResultGroups := make([][]engine.SearchResult, len(config.engine))
|
||||
var finished int
|
||||
var searching int
|
||||
for i := 0; i < len(config.engine); i++ {
|
||||
configEngine := config.engine[i]
|
||||
|
||||
// Check if engine needs to execute for any of the search contexts
|
||||
var needsExecute bool
|
||||
for _, searchCtx := range searchContexts {
|
||||
if configEngine.NeedExectue(searchCtx) {
|
||||
needsExecute = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !needsExecute {
|
||||
continue
|
||||
}
|
||||
|
||||
// Process all search contexts for this engine
|
||||
for _, searchCtx := range searchContexts {
|
||||
if !configEngine.NeedExectue(searchCtx) {
|
||||
continue
|
||||
}
|
||||
args := configEngine.CallArgs(searchCtx)
|
||||
index := i
|
||||
err := configEngine.Client().Call(args.Method, args.Url, args.Headers, args.Body,
|
||||
func(statusCode int, responseHeaders http.Header, responseBody []byte) {
|
||||
defer func() {
|
||||
finished++
|
||||
if finished == searching {
|
||||
// Merge search results from all engines with deduplication
|
||||
var mergedResults []engine.SearchResult
|
||||
seenLinks := make(map[string]bool)
|
||||
for _, results := range searchResultGroups {
|
||||
for _, result := range results {
|
||||
if !seenLinks[result.Link] {
|
||||
seenLinks[result.Link] = true
|
||||
mergedResults = append(mergedResults, result)
|
||||
}
|
||||
}
|
||||
}
|
||||
// Format search results for prompt template
|
||||
var formattedResults []string
|
||||
var formattedReferences []string
|
||||
for j, result := range mergedResults {
|
||||
if config.needReference {
|
||||
formattedResults = append(formattedResults,
|
||||
fmt.Sprintf("[webpage %d begin]\n%s\n[webpage %d end]", j+1, result.Content, j+1))
|
||||
formattedReferences = append(formattedReferences,
|
||||
fmt.Sprintf("[%d] [%s](%s)", j+1, result.Title, result.Link))
|
||||
} else {
|
||||
formattedResults = append(formattedResults,
|
||||
fmt.Sprintf("[webpage begin]\n%s\n[webpage end]", result.Content))
|
||||
}
|
||||
}
|
||||
// Prepare template variables
|
||||
curDate := time.Now().In(time.FixedZone("CST", 8*3600)).Format("2006年1月2日")
|
||||
searchResults := strings.Join(formattedResults, "\n")
|
||||
log.Debugf("searchResults: %s", searchResults)
|
||||
// Fill prompt template
|
||||
prompt := strings.Replace(config.promptTemplate, "{search_results}", searchResults, 1)
|
||||
prompt = strings.Replace(prompt, "{question}", searchContexts[0].Querys[0], 1)
|
||||
prompt = strings.Replace(prompt, "{cur_date}", curDate, 1)
|
||||
// Update request body with processed prompt
|
||||
modifiedBody, err := sjson.SetBytes(body, fmt.Sprintf("messages.%d.content", queryIndex), prompt)
|
||||
if err != nil {
|
||||
log.Errorf("modify request message content failed, err:%v, body:%s", err, body)
|
||||
} else {
|
||||
log.Debugf("modifeid body:%s", modifiedBody)
|
||||
proxywasm.ReplaceHttpRequestBody(modifiedBody)
|
||||
if config.needReference {
|
||||
ctx.SetContext("References", strings.Join(formattedReferences, "\n"))
|
||||
}
|
||||
}
|
||||
proxywasm.ResumeHttpRequest()
|
||||
}
|
||||
}()
|
||||
if statusCode != http.StatusOK {
|
||||
log.Errorf("search call failed, status: %d, engine: %#v", statusCode, configEngine)
|
||||
return
|
||||
}
|
||||
// Append results to existing slice for this engine
|
||||
searchResultGroups[index] = append(searchResultGroups[index], configEngine.ParseResult(searchCtx, responseBody)...)
|
||||
}, args.TimeoutMillisecond)
|
||||
if err != nil {
|
||||
log.Errorf("search call failed, engine: %#v", configEngine)
|
||||
continue
|
||||
}
|
||||
searching++
|
||||
}
|
||||
}
|
||||
if searching > 0 {
|
||||
return types.ActionPause
|
||||
}
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
func onHttpResponseHeaders(ctx wrapper.HttpContext, config Config, log wrapper.Log) types.Action {
|
||||
if !config.needReference {
|
||||
ctx.DontReadResponseBody()
|
||||
return types.ActionContinue
|
||||
}
|
||||
proxywasm.RemoveHttpResponseHeader("content-length")
|
||||
contentType, err := proxywasm.GetHttpResponseHeader("Content-Type")
|
||||
if err != nil || !strings.HasPrefix(contentType, "text/event-stream") {
|
||||
if err != nil {
|
||||
log.Errorf("unable to load content-type header from response: %v", err)
|
||||
}
|
||||
ctx.BufferResponseBody()
|
||||
ctx.SetResponseBodyBufferLimit(DEFAULT_MAX_BODY_BYTES)
|
||||
}
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
func onHttpResponseBody(ctx wrapper.HttpContext, config Config, body []byte, log wrapper.Log) types.Action {
|
||||
references := ctx.GetStringContext("References", "")
|
||||
if references == "" {
|
||||
return types.ActionContinue
|
||||
}
|
||||
content := gjson.GetBytes(body, "choices.0.message.content")
|
||||
modifiedContent := fmt.Sprintf("%s\n\n%s", fmt.Sprintf(config.referenceFormat, references), content)
|
||||
body, err := sjson.SetBytes(body, "choices.0.message.content", modifiedContent)
|
||||
if err != nil {
|
||||
log.Errorf("modify response message content failed, err:%v, body:%s", err, body)
|
||||
return types.ActionContinue
|
||||
}
|
||||
proxywasm.ReplaceHttpResponseBody(body)
|
||||
return types.ActionContinue
|
||||
}
|
||||
|
||||
func onStreamingResponseBody(ctx wrapper.HttpContext, config Config, chunk []byte, isLastChunk bool, log wrapper.Log) []byte {
|
||||
if ctx.GetBoolContext("ReferenceAppended", false) {
|
||||
return chunk
|
||||
}
|
||||
references := ctx.GetStringContext("References", "")
|
||||
if references == "" {
|
||||
return chunk
|
||||
}
|
||||
modifiedChunk, responseReady := setReferencesToFirstMessage(ctx, chunk, fmt.Sprintf(config.referenceFormat, references), log)
|
||||
if responseReady {
|
||||
ctx.SetContext("ReferenceAppended", true)
|
||||
return modifiedChunk
|
||||
} else {
|
||||
return []byte("")
|
||||
}
|
||||
}
|
||||
|
||||
const PARTIAL_MESSAGE_CONTEXT_KEY = "partialMessage"
|
||||
|
||||
func setReferencesToFirstMessage(ctx wrapper.HttpContext, chunk []byte, references string, log wrapper.Log) ([]byte, bool) {
|
||||
if len(chunk) == 0 {
|
||||
log.Debugf("chunk is empty")
|
||||
return nil, false
|
||||
}
|
||||
|
||||
var partialMessage []byte
|
||||
partialMessageI := ctx.GetContext(PARTIAL_MESSAGE_CONTEXT_KEY)
|
||||
if partialMessageI != nil {
|
||||
if pMsg, ok := partialMessageI.([]byte); ok {
|
||||
partialMessage = append(pMsg, chunk...)
|
||||
} else {
|
||||
log.Warnf("invalid partial message type: %T", partialMessageI)
|
||||
partialMessage = chunk
|
||||
}
|
||||
} else {
|
||||
partialMessage = chunk
|
||||
}
|
||||
|
||||
if len(partialMessage) == 0 {
|
||||
log.Debugf("partial message is empty")
|
||||
return nil, false
|
||||
}
|
||||
messages := strings.Split(string(partialMessage), "\n\n")
|
||||
if len(messages) > 1 {
|
||||
firstMessage := messages[0]
|
||||
log.Debugf("first message: %s", firstMessage)
|
||||
firstMessage = strings.TrimPrefix(firstMessage, "data: ")
|
||||
firstMessage = strings.TrimSuffix(firstMessage, "\n")
|
||||
deltaContent := gjson.Get(firstMessage, "choices.0.delta.content")
|
||||
modifiedMessage, err := sjson.Set(firstMessage, "choices.0.delta.content", fmt.Sprintf("%s\n\n%s", references, deltaContent))
|
||||
if err != nil {
|
||||
log.Errorf("modify response delta content failed, err:%v", err)
|
||||
return partialMessage, true
|
||||
}
|
||||
modifiedMessage = fmt.Sprintf("data: %s", modifiedMessage)
|
||||
log.Debugf("modified message: %s", firstMessage)
|
||||
messages[0] = string(modifiedMessage)
|
||||
return []byte(strings.Join(messages, "\n\n")), true
|
||||
}
|
||||
ctx.SetContext(PARTIAL_MESSAGE_CONTEXT_KEY, partialMessage)
|
||||
return nil, false
|
||||
}
|
||||
214
plugins/wasm-go/extensions/ai-search/prompts/arxiv.md
Normal file
214
plugins/wasm-go/extensions/ai-search/prompts/arxiv.md
Normal file
@@ -0,0 +1,214 @@
|
||||
# 目标
|
||||
你需要分析**用户发送的消息**，是否需要查询搜索引擎(Google/Bing)/论文资料库(Arxiv)，并按照如下情况回复相应内容:
|
||||
|
||||
## 情况一：不需要查询搜索引擎/论文资料/私有知识库
|
||||
### 情况举例：
|
||||
1. **用户发送的消息**不是在提问或寻求帮助
|
||||
2. **用户发送的消息**是要求翻译文字
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，如果符合，则按照下面**回复内容示例**进行回复，注意不要输出思考过程
|
||||
|
||||
### 回复内容示例：
|
||||
none
|
||||
|
||||
## 情况二：需要查询搜索引擎/论文资料
|
||||
### 情况举例：
|
||||
1. 答复**用户发送的消息**，需依赖互联网上最新的资料
|
||||
2. 答复**用户发送的消息**，需依赖论文等专业资料
|
||||
3. 通过查询资料，可以更好地答复**用户发送的消息**
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，以及其他需要查询资料的情况，如果符合，按照以下步骤思考，并按照下面**回复内容示例**进行回复，注意不要输出思考过程:
|
||||
1. What: 分析要答复**用户发送的消息**，需要了解什么知识和资料
|
||||
2. Where: 判断了解这个知识和资料要向Google等搜索引擎提问，还是向Arxiv论文资料库进行查询，或者需要同时查询多个地方
|
||||
3. How: 分析对于要查询的知识和资料，应该提出什么样的问题
|
||||
4. Adjust: 明确要向什么地方查询什么问题后，按下面方式对问题进行调整
|
||||
4.1. 向搜索引擎提问：用一句话概括问题，并且针对搜索引擎做问题优化
|
||||
4.2. 向Arxiv论文资料库提问：
|
||||
4.2.1. 明确问题所属领域，然后确定Arxiv的Category值，Category可选的枚举如下:
|
||||
- cs.AI: Artificial Intelligence
|
||||
- cs.AR: Hardware Architecture
|
||||
- cs.CC: Computational Complexity
|
||||
- cs.CE: Computational Engineering, Finance, and Science
|
||||
- cs.CG: Computational Geometry
|
||||
- cs.CL: Computation and Language
|
||||
- cs.CR: Cryptography and Security
|
||||
- cs.CV: Computer Vision and Pattern Recognition
|
||||
- cs.CY: Computers and Society
|
||||
- cs.DB: Databases则按照下面**回复内容**进行回复
|
||||
- cs.DC: Distributed, Parallel, and Cluster Computing
|
||||
- cs.DL: Digital Libraries
|
||||
- cs.DM: Discrete Mathematics
|
||||
- cs.DS: Data Structures and Algorithms
|
||||
- cs.ET: Emerging Technologies
|
||||
- cs.FL: Formal Languages and Automata Theory
|
||||
- cs.GL: General Literature
|
||||
- cs.GR: Graphics
|
||||
- cs.GT: Computer Science and Game Theory
|
||||
- cs.HC: Human-Computer Interaction
|
||||
- cs.IR: Information Retrieval
|
||||
- cs.IT: Information Theory
|
||||
- cs.LG: Machine Learning
|
||||
- cs.LO: Logic in Computer Science
|
||||
- cs.MA: Multiagent Systems
|
||||
- cs.MM: Multimedia
|
||||
- cs.MS: Mathematical Software
|
||||
- cs.NA: Numerical Analysis
|
||||
- cs.NE: Neural and Evolutionary Computing
|
||||
- cs.NI: Networking and Internet Architecture
|
||||
- cs.OH: Other Computer Science
|
||||
- cs.OS: Operating Systems
|
||||
- cs.PF: Performance
|
||||
- cs.PL: Programming Languages
|
||||
- cs.RO: Robotics
|
||||
- cs.SC: Symbolic Computation
|
||||
- cs.SD: Sound
|
||||
- cs.SE: Software Engineering
|
||||
- cs.SI: Social and Information Networks
|
||||
- cs.SY: Systems and Control
|
||||
- econ.EM: Econometrics
|
||||
- econ.GN: General Economics
|
||||
- econ.TH: Theoretical Economics
|
||||
- eess.AS: Audio and Speech Processing
|
||||
- eess.IV: Image and Video Processing
|
||||
- eess.SP: Signal Processing
|
||||
- eess.SY: Systems and Control
|
||||
- math.AC: Commutative Algebra
|
||||
- math.AG: Algebraic Geometry
|
||||
- math.AP: Analysis of PDEs
|
||||
- math.AT: Algebraic Topology
|
||||
- math.CA: Classical Analysis and ODEs
|
||||
- math.CO: Combinatorics
|
||||
- math.CT: Category Theory
|
||||
- math.CV: Complex Variables
|
||||
- math.DG: Differential Geometry
|
||||
- math.DS: Dynamical Systems
|
||||
- math.FA: Functional Analysis
|
||||
- math.GM: General Mathematics
|
||||
- math.GN: General Topology
|
||||
- math.GR: Group Theory
|
||||
- math.GT: Geometric Topology
|
||||
- math.HO: History and Overview
|
||||
- math.IT: Information Theory
|
||||
- math.KT: K-Theory and Homology
|
||||
- math.LO: Logic
|
||||
- math.MG: Metric Geometry
|
||||
- math.MP: Mathematical Physics
|
||||
- math.NA: Numerical Analysis
|
||||
- math.NT: Number Theory
|
||||
- math.OA: Operator Algebras
|
||||
- math.OC: Optimization and Control
|
||||
- math.PR: Probability
|
||||
- math.QA: Quantum Algebra
|
||||
- math.RA: Rings and Algebras
|
||||
- math.RT: Representation Theory
|
||||
- math.SG: Symplectic Geometry
|
||||
- math.SP: Spectral Theory
|
||||
- math.ST: Statistics Theory
|
||||
- astro-ph.CO: Cosmology and Nongalactic Astrophysics
|
||||
- astro-ph.EP: Earth and Planetary Astrophysics
|
||||
- astro-ph.GA: Astrophysics of Galaxies
|
||||
- astro-ph.HE: High Energy Astrophysical Phenomena
|
||||
- astro-ph.IM: Instrumentation and Methods for Astrophysics
|
||||
- astro-ph.SR: Solar and Stellar Astrophysics
|
||||
- cond-mat.dis-nn: Disordered Systems and Neural Networks
|
||||
- cond-mat.mes-hall: Mesoscale and Nanoscale Physics
|
||||
- cond-mat.mtrl-sci: Materials Science
|
||||
- cond-mat.other: Other Condensed Matter
|
||||
- cond-mat.quant-gas: Quantum Gases
|
||||
- cond-mat.soft: Soft Condensed Matter
|
||||
- cond-mat.stat-mech: Statistical Mechanics
|
||||
- cond-mat.str-el: Strongly Correlated Electrons
|
||||
- cond-mat.supr-con: Superconductivity
|
||||
- gr-qc: General Relativity and Quantum Cosmology
|
||||
- hep-ex: High Energy Physics - Experiment
|
||||
- hep-lat: High Energy Physics - Lattice
|
||||
- hep-ph: High Energy Physics - Phenomenology
|
||||
- hep-th: High Energy Physics - Theory
|
||||
- math-ph: Mathematical Physics
|
||||
- nlin.AO: Adaptation and Self-Organizing Systems
|
||||
- nlin.CD: Chaotic Dynamics
|
||||
- nlin.CG: Cellular Automata and Lattice Gases
|
||||
- nlin.PS: Pattern Formation and Solitons
|
||||
- nlin.SI: Exactly Solvable and Integrable Systems
|
||||
- nucl-ex: Nuclear Experiment
|
||||
- nucl-th: Nuclear Theory
|
||||
- physics.acc-ph: Accelerator Physics
|
||||
- physics.ao-ph: Atmospheric and Oceanic Physics
|
||||
- physics.app-ph: Applied Physics
|
||||
- physics.atm-clus: Atomic and Molecular Clusters
|
||||
- physics.atom-ph: Atomic Physics
|
||||
- physics.bio-ph: Biological Physics
|
||||
- physics.chem-ph: Chemical Physics
|
||||
- physics.class-ph: Classical Physics
|
||||
- physics.comp-ph: Computational Physics
|
||||
- physics.data-an: Data Analysis, Statistics and Probability
|
||||
- physics.ed-ph: Physics Education
|
||||
- physics.flu-dyn: Fluid Dynamics
|
||||
- physics.gen-ph: General Physics
|
||||
- physics.geo-ph: Geophysics
|
||||
- physics.hist-ph: History and Philosophy of Physics
|
||||
- physics.ins-det: Instrumentation and Detectors
|
||||
- physics.med-ph: Medical Physics
|
||||
- physics.optics: Optics
|
||||
- physics.plasm-ph: Plasma Physics
|
||||
- physics.pop-ph: Popular Physics
|
||||
- physics.soc-ph: Physics and Society
|
||||
- physics.space-ph: Space Physics
|
||||
- quant-ph: Quantum Physics
|
||||
- q-bio.BM: Biomolecules
|
||||
- q-bio.CB: Cell Behavior
|
||||
- q-bio.GN: Genomics
|
||||
- q-bio.MN: Molecular Networks
|
||||
- q-bio.NC: Neurons and Cognition
|
||||
- q-bio.OT: Other Quantitative Biology
|
||||
- q-bio.PE: Populations and Evolution
|
||||
- q-bio.QM: Quantitative Methods
|
||||
- q-bio.SC: Subcellular Processes
|
||||
- q-bio.TO: Tissues and Organs
|
||||
- q-fin.CP: Computational Finance
|
||||
- q-fin.EC: Economics
|
||||
- q-fin.GN: General Finance
|
||||
- q-fin.MF: Mathematical Finance
|
||||
- q-fin.PM: Portfolio Management
|
||||
- q-fin.PR: Pricing of Securities
|
||||
- q-fin.RM: Risk Management
|
||||
- q-fin.ST: Statistical Finance
|
||||
- q-fin.TR: Trading and Market Microstructure
|
||||
- stat.AP: Applications
|
||||
- stat.CO: Computation
|
||||
- stat.ME: Methodology
|
||||
- stat.ML: Machine Learning
|
||||
- stat.OT: Other Statistics
|
||||
- stat.TH: Statistics Theory
|
||||
4.2.2. 根据问题所属领域，将问题拆分成多组关键词的组合，同时组合中的关键词个数尽量不要超过3个
|
||||
5. Final: 按照下面**回复内容示例**进行回复，注意:
|
||||
- 不要输出思考过程
|
||||
- 可以向多个查询目标分别查询多次，多个查询用换行分隔，总查询次数控制在5次以内
|
||||
- 查询搜索引擎时，需要以"internet:"开头
|
||||
- 查询Arxiv论文时，需要以Arxiv的Category值开头，例如"cs.AI:"
|
||||
- 查询Arxiv论文时，优先用英文表述关键词进行搜索
|
||||
- 当用多个关键词查询时，关键词之间用","分隔
|
||||
- 尽量满足**用户发送的消息**中的搜索要求，例如用户要求用英文搜索，则需用英文表述问题和关键词
|
||||
- 用户如果没有要求搜索语言，则用和**用户发送的消息**一致的语言表述问题和关键词
|
||||
- 如果**用户发送的消息**使用中文，至少要有一条向搜索引擎查询的中文问题
|
||||
|
||||
### 回复内容示例：
|
||||
|
||||
#### 用不同语言查询多次搜索引擎
|
||||
internet: 黄金价格走势
|
||||
internet: The trend of gold prices
|
||||
|
||||
#### 向Arxiv的多个类目查询多次
|
||||
cs.AI: attention mechanism
|
||||
cs.AI: neuron
|
||||
q-bio.NC: brain,attention mechanism
|
||||
|
||||
#### 向多个查询目标查询多次
|
||||
internet: 中国未来房价趋势
|
||||
internet: 最新中国经济政策
|
||||
econ.TH: policy, real estate
|
||||
|
||||
# 用户发送的消息为：
|
||||
{question}
|
||||
221
plugins/wasm-go/extensions/ai-search/prompts/full.md
Normal file
221
plugins/wasm-go/extensions/ai-search/prompts/full.md
Normal file
@@ -0,0 +1,221 @@
|
||||
# 目标
|
||||
你需要分析**用户发送的消息**，是否需要查询搜索引擎(Google/Bing)/论文资料库(Arxiv)/私有知识库，并按照如下情况回复相应内容:
|
||||
|
||||
## 情况一：不需要查询搜索引擎/论文资料/私有知识库
|
||||
### 情况举例：
|
||||
1. **用户发送的消息**不是在提问或寻求帮助
|
||||
2. **用户发送的消息**是要求翻译文字
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，如果符合，则按照下面**回复内容示例**进行回复，注意不要输出思考过程
|
||||
|
||||
### 回复内容示例：
|
||||
none
|
||||
|
||||
## 情况二：需要查询搜索引擎/论文资料/私有知识库
|
||||
### 情况举例：
|
||||
1. 答复**用户发送的消息**，需依赖互联网上最新的资料
|
||||
2. 答复**用户发送的消息**，需依赖论文等专业资料
|
||||
3. 通过查询资料，可以更好地答复**用户发送的消息**
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，以及其他需要查询资料的情况，如果符合，按照以下步骤思考，并按照下面**回复内容示例**进行回复，注意不要输出思考过程:
|
||||
1. What: 分析要答复**用户发送的消息**，需要了解什么知识和资料
|
||||
2. Where: 判断了解这个知识和资料要向Google等搜索引擎提问，还是向Arxiv论文资料库进行查询，还是向私有知识库进行查询，或者需要同时查询多个地方
|
||||
3. How: 分析对于要查询的知识和资料，应该提出什么样的问题
|
||||
4. Adjust: 明确要向什么地方查询什么问题后，按下面方式对问题进行调整
|
||||
4.1. 向搜索引擎提问：用一句话概括问题，并且针对搜索引擎做问题优化
|
||||
4.2. 向私有知识库提问：将问题拆分成多组关键词的组合，同时组合中的关键词个数尽量不要超过3个
|
||||
4.3. 向Arxiv论文资料库提问：
|
||||
4.3.1. 明确问题所属领域，然后确定Arxiv的Category值，Category可选的枚举如下:
|
||||
- cs.AI: Artificial Intelligence
|
||||
- cs.AR: Hardware Architecture
|
||||
- cs.CC: Computational Complexity
|
||||
- cs.CE: Computational Engineering, Finance, and Science
|
||||
- cs.CG: Computational Geometry
|
||||
- cs.CL: Computation and Language
|
||||
- cs.CR: Cryptography and Security
|
||||
- cs.CV: Computer Vision and Pattern Recognition
|
||||
- cs.CY: Computers and Society
|
||||
- cs.DB: Databases则按照下面**回复内容**进行回复
|
||||
- cs.DC: Distributed, Parallel, and Cluster Computing
|
||||
- cs.DL: Digital Libraries
|
||||
- cs.DM: Discrete Mathematics
|
||||
- cs.DS: Data Structures and Algorithms
|
||||
- cs.ET: Emerging Technologies
|
||||
- cs.FL: Formal Languages and Automata Theory
|
||||
- cs.GL: General Literature
|
||||
- cs.GR: Graphics
|
||||
- cs.GT: Computer Science and Game Theory
|
||||
- cs.HC: Human-Computer Interaction
|
||||
- cs.IR: Information Retrieval
|
||||
- cs.IT: Information Theory
|
||||
- cs.LG: Machine Learning
|
||||
- cs.LO: Logic in Computer Science
|
||||
- cs.MA: Multiagent Systems
|
||||
- cs.MM: Multimedia
|
||||
- cs.MS: Mathematical Software
|
||||
- cs.NA: Numerical Analysis
|
||||
- cs.NE: Neural and Evolutionary Computing
|
||||
- cs.NI: Networking and Internet Architecture
|
||||
- cs.OH: Other Computer Science
|
||||
- cs.OS: Operating Systems
|
||||
- cs.PF: Performance
|
||||
- cs.PL: Programming Languages
|
||||
- cs.RO: Robotics
|
||||
- cs.SC: Symbolic Computation
|
||||
- cs.SD: Sound
|
||||
- cs.SE: Software Engineering
|
||||
- cs.SI: Social and Information Networks
|
||||
- cs.SY: Systems and Control
|
||||
- econ.EM: Econometrics
|
||||
- econ.GN: General Economics
|
||||
- econ.TH: Theoretical Economics
|
||||
- eess.AS: Audio and Speech Processing
|
||||
- eess.IV: Image and Video Processing
|
||||
- eess.SP: Signal Processing
|
||||
- eess.SY: Systems and Control
|
||||
- math.AC: Commutative Algebra
|
||||
- math.AG: Algebraic Geometry
|
||||
- math.AP: Analysis of PDEs
|
||||
- math.AT: Algebraic Topology
|
||||
- math.CA: Classical Analysis and ODEs
|
||||
- math.CO: Combinatorics
|
||||
- math.CT: Category Theory
|
||||
- math.CV: Complex Variables
|
||||
- math.DG: Differential Geometry
|
||||
- math.DS: Dynamical Systems
|
||||
- math.FA: Functional Analysis
|
||||
- math.GM: General Mathematics
|
||||
- math.GN: General Topology
|
||||
- math.GR: Group Theory
|
||||
- math.GT: Geometric Topology
|
||||
- math.HO: History and Overview
|
||||
- math.IT: Information Theory
|
||||
- math.KT: K-Theory and Homology
|
||||
- math.LO: Logic
|
||||
- math.MG: Metric Geometry
|
||||
- math.MP: Mathematical Physics
|
||||
- math.NA: Numerical Analysis
|
||||
- math.NT: Number Theory
|
||||
- math.OA: Operator Algebras
|
||||
- math.OC: Optimization and Control
|
||||
- math.PR: Probability
|
||||
- math.QA: Quantum Algebra
|
||||
- math.RA: Rings and Algebras
|
||||
- math.RT: Representation Theory
|
||||
- math.SG: Symplectic Geometry
|
||||
- math.SP: Spectral Theory
|
||||
- math.ST: Statistics Theory
|
||||
- astro-ph.CO: Cosmology and Nongalactic Astrophysics
|
||||
- astro-ph.EP: Earth and Planetary Astrophysics
|
||||
- astro-ph.GA: Astrophysics of Galaxies
|
||||
- astro-ph.HE: High Energy Astrophysical Phenomena
|
||||
- astro-ph.IM: Instrumentation and Methods for Astrophysics
|
||||
- astro-ph.SR: Solar and Stellar Astrophysics
|
||||
- cond-mat.dis-nn: Disordered Systems and Neural Networks
|
||||
- cond-mat.mes-hall: Mesoscale and Nanoscale Physics
|
||||
- cond-mat.mtrl-sci: Materials Science
|
||||
- cond-mat.other: Other Condensed Matter
|
||||
- cond-mat.quant-gas: Quantum Gases
|
||||
- cond-mat.soft: Soft Condensed Matter
|
||||
- cond-mat.stat-mech: Statistical Mechanics
|
||||
- cond-mat.str-el: Strongly Correlated Electrons
|
||||
- cond-mat.supr-con: Superconductivity
|
||||
- gr-qc: General Relativity and Quantum Cosmology
|
||||
- hep-ex: High Energy Physics - Experiment
|
||||
- hep-lat: High Energy Physics - Lattice
|
||||
- hep-ph: High Energy Physics - Phenomenology
|
||||
- hep-th: High Energy Physics - Theory
|
||||
- math-ph: Mathematical Physics
|
||||
- nlin.AO: Adaptation and Self-Organizing Systems
|
||||
- nlin.CD: Chaotic Dynamics
|
||||
- nlin.CG: Cellular Automata and Lattice Gases
|
||||
- nlin.PS: Pattern Formation and Solitons
|
||||
- nlin.SI: Exactly Solvable and Integrable Systems
|
||||
- nucl-ex: Nuclear Experiment
|
||||
- nucl-th: Nuclear Theory
|
||||
- physics.acc-ph: Accelerator Physics
|
||||
- physics.ao-ph: Atmospheric and Oceanic Physics
|
||||
- physics.app-ph: Applied Physics
|
||||
- physics.atm-clus: Atomic and Molecular Clusters
|
||||
- physics.atom-ph: Atomic Physics
|
||||
- physics.bio-ph: Biological Physics
|
||||
- physics.chem-ph: Chemical Physics
|
||||
- physics.class-ph: Classical Physics
|
||||
- physics.comp-ph: Computational Physics
|
||||
- physics.data-an: Data Analysis, Statistics and Probability
|
||||
- physics.ed-ph: Physics Education
|
||||
- physics.flu-dyn: Fluid Dynamics
|
||||
- physics.gen-ph: General Physics
|
||||
- physics.geo-ph: Geophysics
|
||||
- physics.hist-ph: History and Philosophy of Physics
|
||||
- physics.ins-det: Instrumentation and Detectors
|
||||
- physics.med-ph: Medical Physics
|
||||
- physics.optics: Optics
|
||||
- physics.plasm-ph: Plasma Physics
|
||||
- physics.pop-ph: Popular Physics
|
||||
- physics.soc-ph: Physics and Society
|
||||
- physics.space-ph: Space Physics
|
||||
- quant-ph: Quantum Physics
|
||||
- q-bio.BM: Biomolecules
|
||||
- q-bio.CB: Cell Behavior
|
||||
- q-bio.GN: Genomics
|
||||
- q-bio.MN: Molecular Networks
|
||||
- q-bio.NC: Neurons and Cognition
|
||||
- q-bio.OT: Other Quantitative Biology
|
||||
- q-bio.PE: Populations and Evolution
|
||||
- q-bio.QM: Quantitative Methods
|
||||
- q-bio.SC: Subcellular Processes
|
||||
- q-bio.TO: Tissues and Organs
|
||||
- q-fin.CP: Computational Finance
|
||||
- q-fin.EC: Economics
|
||||
- q-fin.GN: General Finance
|
||||
- q-fin.MF: Mathematical Finance
|
||||
- q-fin.PM: Portfolio Management
|
||||
- q-fin.PR: Pricing of Securities
|
||||
- q-fin.RM: Risk Management
|
||||
- q-fin.ST: Statistical Finance
|
||||
- q-fin.TR: Trading and Market Microstructure
|
||||
- stat.AP: Applications
|
||||
- stat.CO: Computation
|
||||
- stat.ME: Methodology
|
||||
- stat.ML: Machine Learning
|
||||
- stat.OT: Other Statistics
|
||||
- stat.TH: Statistics Theory
|
||||
4.3.2. 根据问题所属领域，将问题拆分成多组关键词的组合，同时组合中的关键词个数尽量不要超过3个
|
||||
5. Final: 按照下面**回复内容示例**进行回复，注意:
|
||||
- 不要输出思考过程
|
||||
- 可以向多个查询目标分别查询多次，多个查询用换行分隔，总查询次数控制在5次以内
|
||||
- 查询搜索引擎时，需要以"internet:"开头
|
||||
- 查询私有知识库时，需要以"private:"开头
|
||||
- 查询Arxiv论文时，需要以Arxiv的Category值开头，例如"cs.AI:"
|
||||
- 查询Arxiv论文时，优先用英文表述关键词进行搜索
|
||||
- 当用多个关键词查询时，关键词之间用","分隔
|
||||
- 尽量满足**用户发送的消息**中的搜索要求，例如用户要求用英文搜索，则需用英文表述问题和关键词
|
||||
- 用户如果没有要求搜索语言，则用和**用户发送的消息**一致的语言表述问题和关键词
|
||||
- 如果**用户发送的消息**使用中文，至少要有一条向搜索引擎查询的中文问题
|
||||
|
||||
### 回复内容示例：
|
||||
|
||||
#### 用不同语言查询多次搜索引擎
|
||||
internet: 黄金价格走势
|
||||
internet: The trend of gold prices
|
||||
|
||||
#### 向Arxiv的多个类目查询多次
|
||||
cs.AI: attention mechanism
|
||||
cs.AI: neuron
|
||||
q-bio.NC: brain,attention mechanism
|
||||
|
||||
#### 向私有知识库查询多次
|
||||
private: 电子钱包,密码
|
||||
private: 张三,身份证号
|
||||
|
||||
#### 向多个查询目标查询多次
|
||||
internet: 中国未来房价趋势
|
||||
internet: 最新中国经济政策
|
||||
econ.TH: policy, real estate
|
||||
private: 财务状况
|
||||
|
||||
# 用户发送的消息为：
|
||||
{question}
|
||||
41
plugins/wasm-go/extensions/ai-search/prompts/internet.md
Normal file
41
plugins/wasm-go/extensions/ai-search/prompts/internet.md
Normal file
@@ -0,0 +1,41 @@
|
||||
# 目标
|
||||
你需要分析**用户发送的消息**，是否需要查询搜索引擎(Google/Bing)，并按照如下情况回复相应内容:
|
||||
|
||||
## 情况一：不需要查询搜索引擎
|
||||
### 情况举例：
|
||||
1. **用户发送的消息**不是在提问或寻求帮助
|
||||
2. **用户发送的消息**是要求翻译文字
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，如果符合，则按照下面**回复内容示例**进行回复，注意不要输出思考过程
|
||||
|
||||
### 回复内容示例：
|
||||
none
|
||||
|
||||
## 情况二：需要查询搜索引擎
|
||||
### 情况举例：
|
||||
1. 答复**用户发送的消息**，需依赖互联网上最新的资料
|
||||
2. 答复**用户发送的消息**，需依赖论文等专业资料
|
||||
3. 通过查询资料，可以更好地答复**用户发送的消息**
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，以及其他需要查询资料的情况，如果符合，按照以下步骤思考，并按照下面**回复内容示例**进行回复，注意不要输出思考过程:
|
||||
1. What: 分析要答复**用户发送的消息**，需要了解什么知识和资料
|
||||
2. How: 分析对于要查询的知识和资料，应该提出什么样的问题
|
||||
3. Adjust: 明确查询什么问题后，用一句话概括问题，并且针对搜索引擎做问题优化
|
||||
4. Final: 按照下面**回复内容示例**进行回复，注意:
|
||||
- 不要输出思考过程
|
||||
- 可以查询多次，多个查询用换行分隔，总查询次数控制在5次以内
|
||||
- 需要以"internet:"开头
|
||||
- 尽量满足**用户发送的消息**中的搜索要求，例如用户要求用英文搜索，则需用英文表述问题和关键词
|
||||
- 用户如果没有要求搜索语言，则用和**用户发送的消息**一致的语言表述问题和关键词
|
||||
- 如果**用户发送的消息**使用中文，至少要有一条向搜索引擎查询的中文问题
|
||||
|
||||
### 回复内容示例：
|
||||
|
||||
#### 用不同语言查询多次搜索引擎
|
||||
internet: 黄金价格走势
|
||||
internet: The trend of gold prices
|
||||
|
||||
# 用户发送的消息为：
|
||||
{question}
|
||||
55
plugins/wasm-go/extensions/ai-search/prompts/private.md
Normal file
55
plugins/wasm-go/extensions/ai-search/prompts/private.md
Normal file
@@ -0,0 +1,55 @@
|
||||
# 目标
|
||||
你需要分析**用户发送的消息**，是否需要查询搜索引擎(Google/Bing)/私有知识库，并按照如下情况回复相应内容:
|
||||
|
||||
## 情况一：不需要查询搜索引擎/私有知识库
|
||||
### 情况举例：
|
||||
1. **用户发送的消息**不是在提问或寻求帮助
|
||||
2. **用户发送的消息**是要求翻译文字
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，如果符合，则按照下面**回复内容示例**进行回复，注意不要输出思考过程
|
||||
|
||||
### 回复内容示例：
|
||||
none
|
||||
|
||||
## 情况二：需要查询搜索引擎/私有知识库
|
||||
### 情况举例：
|
||||
1. 答复**用户发送的消息**，需依赖互联网上最新的资料
|
||||
2. 答复**用户发送的消息**，需依赖论文等专业资料
|
||||
3. 通过查询资料，可以更好地答复**用户发送的消息**
|
||||
|
||||
### 思考过程
|
||||
根据上面的**情况举例**，以及其他需要查询资料的情况，如果符合，按照以下步骤思考，并按照下面**回复内容示例**进行回复，注意不要输出思考过程:
|
||||
1. What: 分析要答复**用户发送的消息**，需要了解什么知识和资料
|
||||
2. Where: 判断了解这个知识和资料要向Google等搜索引擎提问，还是向私有知识库进行查询，或者需要同时查询多个地方
|
||||
3. How: 分析对于要查询的知识和资料，应该提出什么样的问题
|
||||
4. Adjust: 明确要向什么地方查询什么问题后，按下面方式对问题进行调整
|
||||
4.1. 向搜索引擎提问：用一句话概括问题，并且针对搜索引擎做问题优化
|
||||
4.2. 向私有知识库提问：将问题拆分成多组关键词的组合，同时组合中的关键词个数尽量不要超过3个
|
||||
5. Final: 按照下面**回复内容示例**进行回复，注意:
|
||||
- 不要输出思考过程
|
||||
- 可以向多个查询目标分别查询多次，多个查询用换行分隔，总查询次数控制在5次以内
|
||||
- 查询搜索引擎时，需要以"internet:"开头
|
||||
- 查询私有知识库时，需要以"private:"开头
|
||||
- 当用多个关键词查询时，关键词之间用","分隔
|
||||
- 尽量满足**用户发送的消息**中的搜索要求，例如用户要求用英文搜索，则需用英文表述问题和关键词
|
||||
- 用户如果没有要求搜索语言，则用和**用户发送的消息**一致的语言表述问题和关键词
|
||||
- 如果**用户发送的消息**使用中文，至少要有一条向搜索引擎查询的中文问题
|
||||
|
||||
### 回复内容示例：
|
||||
|
||||
#### 用不同语言查询多次搜索引擎
|
||||
internet: 黄金价格走势
|
||||
internet: The trend of gold prices
|
||||
|
||||
#### 向私有知识库查询多次
|
||||
private: 电子钱包,密码
|
||||
private: 张三,身份证号
|
||||
|
||||
#### 向多个查询目标查询多次
|
||||
internet: 中国未来房价趋势
|
||||
internet: 最新中国经济政策
|
||||
private: 财务状况
|
||||
|
||||
# 用户发送的消息为：
|
||||
{question}
|
||||
@@ -0,0 +1,56 @@
|
||||
import argparse
|
||||
import requests
|
||||
import time
|
||||
import json
|
||||
|
||||
def main():
|
||||
# 解析命令行参数
|
||||
parser = argparse.ArgumentParser(description='AI Search Test Script')
|
||||
parser.add_argument('--question', required=True, help='The question to analyze')
|
||||
parser.add_argument('--prompt', required=True, help='The prompt file to analyze')
|
||||
args = parser.parse_args()
|
||||
|
||||
# 读取并解析prompts.md模板
|
||||
# 这里假设prompts.md已经复制到当前目录
|
||||
with open(args.prompt, 'r', encoding='utf-8') as f:
|
||||
prompt_template = f.read()
|
||||
|
||||
# 替换模板中的{question}变量
|
||||
prompt = prompt_template.replace('{question}', args.question)
|
||||
|
||||
# 准备请求数据
|
||||
headers = {
|
||||
'Content-Type': 'application/json',
|
||||
}
|
||||
data = {
|
||||
"model": "deepseek-v3",
|
||||
"max_tokens": 100,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": prompt
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
# 发送请求并计时
|
||||
start_time = time.time()
|
||||
try:
|
||||
response = requests.post(
|
||||
'http://localhost:8080/v1/chat/completions',
|
||||
headers=headers,
|
||||
data=json.dumps(data)
|
||||
)
|
||||
response.raise_for_status()
|
||||
end_time = time.time()
|
||||
|
||||
# 处理响应
|
||||
result = response.json()
|
||||
print("Response:")
|
||||
print(result['choices'][0]['message']['content'])
|
||||
print(f"\nRequest took {end_time - start_time:.2f} seconds")
|
||||
except requests.exceptions.RequestException as e:
|
||||
print(f"Request failed: {e}")
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
Reference in New Issue
Block a user