feat: ai-rag plugin #11568

shreemaan-abhishek · 2024-09-11T11:18:47Z

Description

Implement the ai-rag plugin that uses RAG (Retrieval Augmented Generation), to return information about untrained events/facts from LLMs.

Checklist

I have explained the need for this PR and the problem it solves
I have explained the changes or the new features added to this PR
I have added tests corresponding to this change
I have updated the documentation to reflect this change
I have verified that this change is backward compatible (If not, please discuss on the APISIX mailing list first)

apisix/plugins/ai-rag.lua

shreemaan-abhishek · 2024-09-11T11:25:43Z

apisix/plugins/ai-rag.lua

+
+    request_schema.properties.ai_rag.properties.vector_search = vs_req_schema
+    request_schema.properties.ai_rag.properties.embeddings = emb_req_schema
+


construct a request_body_schema according to the service providers

apisix/plugins/ai-rag.lua

shreemaan-abhishek · 2024-09-11T11:27:21Z

apisix/plugins/ai-rag.lua

+    core.log.error("dibag err: ", err)
+    core.log.warn("dibag res: ", core.json.encode(res, true))
+
+    body_tab["ai_rag"] = nil


remove ai_rag from request body because their purpose is served and also these values will cause failure when proxying requests to LLM.

comment the purpose in source code?

shreemaan-abhishek · 2024-09-11T11:29:38Z

apisix/plugins/ai-rag.lua

+    if not body_tab.messages then
+        body_tab.messages = {}
+    end
+    decorate(decorator_conf, body_tab)


we need to take the vector search result and append it to the messages field in request body. This messages will be sent to LLM to get the response.

Based on the code, I think you mean "prepend" instead of "append"?

ah, yes. My bad, it should be prepend.

zhoujiexiong

Is documentation related to this PR being prepared?

zhoujiexiong · 2024-09-21T12:46:19Z

t/plugin/ai-rag.t

+
+                    if ngx.req.get_method() ~= "POST" then
+                        ngx.status = 400
+                        ngx.say("Unsupported request method: ", ngx.req.get_method())


then return?

zhoujiexiong · 2024-09-21T12:50:19Z

t/plugin/ai-rag.t

+                        return
+                    end
+
+                    if header_auth == "key" then


Repeated judgment?

zhoujiexiong · 2024-09-21T12:59:16Z

apisix/plugins/ai-rag/embeddings/azure_openai.lua

+            ["Content-Type"] = "application/json",
+            ["api-key"] = conf.api_key,
+        },
+        body = core.json.encode(body)


zhoujiexiong · 2024-09-21T13:00:12Z

apisix/plugins/ai-rag/embeddings/azure_openai.lua

+        body = core.json.encode(body)
+    })
+    if not res or not res.body then
+        return nil, err


Suggested change

return nil, err

return nil, STATUS, err

zhoujiexiong · 2024-09-21T13:03:47Z

apisix/plugins/ai-rag/embeddings/azure_openai.lua

+        return nil, internal_server_error, err
+    end
+
+    if type(res_tab.data) ~= "table" or #res_tab.data < 1 then


Suggested change

if type(res_tab.data) ~= "table" or #res_tab.data < 1 then

if type(res_tab.data) ~= "table" or core.table.isempty(res_tab.data) then

zhoujiexiong · 2024-09-21T13:27:05Z

apisix/plugins/ai-rag.lua

+    end
+
+    local embeddings_provider = next(conf.embeddings_provider)
+    local embeddings_provider_conf = conf.embeddings_provider[next(conf.embeddings_provider)]


Suggested change

local embeddings_provider_conf = conf.embeddings_provider[next(conf.embeddings_provider)]

local embeddings_provider_conf = conf.embeddings_provider[embeddings_provider]

zhoujiexiong · 2024-09-21T13:31:22Z

apisix/plugins/ai-rag.lua

+    local vs_req_schema = vector_search_driver.request_schema
+    local emb_req_schema = embeddings_driver.request_schema
+
+    request_schema.properties.ai_rag.properties.vector_search = vs_req_schema


Are there risks with module variables in concurrent scenarios?

I don't think so, can you elaborate?

I checked the process again and there is no risk of race condition. Ignore it.

zhoujiexiong · 2024-09-21T13:35:52Z

apisix/plugins/ai-rag.lua

+    core.log.error("dibag err: ", err)
+    core.log.warn("dibag res: ", core.json.encode(res, true))
+
+    body_tab["ai_rag"] = nil


comment the purpose in source code?

zhoujiexiong · 2024-09-21T13:39:32Z

apisix/plugins/ai-rag.lua

+    if not body_tab.messages then
+        body_tab.messages = {}
+    end
+    decorate(decorator_conf, body_tab)


Based on the code, I think you mean "prepend" instead of "append"?

zhoujiexiong · 2024-09-21T13:43:57Z

apisix/plugins/ai-rag/embeddings/azure_openai.lua

+-- limitations under the License.
+--
+local core = require("apisix.core")
+local internal_server_error = ngx.HTTP_INTERNAL_SERVER_ERROR


Suggested change

local internal_server_error = ngx.HTTP_INTERNAL_SERVER_ERROR

local INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR

const var, sugg. this style, others elsewhere the same :)

bzp2010

There is no documentation and I don't know how to use it.

bzp2010 · 2024-09-22T19:05:41Z

apisix/plugins/ai-rag.lua

+local azure_ai_search_schema = {
+    type = "object",
+    properties = {
+        endpoint = {
+            type = "string",
+        },
+        api_key = {
+            type = "string",
+        },
+    }
+}
+
+local azure_openai_embeddings = {
+    type = "object",
+    properties = {
+        endpoint = {
+            type = "string",
+        },
+        api_key = {
+            type = "string",
+        },
+    },
+    required = { "endpoint", "api_key" }
+}


This should be moved to the provider Lua file.

membphis · 2024-09-25T03:25:14Z

apisix/plugins/ai-prompt-decorator.lua

@@ -114,4 +114,5 @@ function _M.rewrite(conf, ctx)
 end


+_M.__decorate = decorate -- for ai-rag plugin


it is unsafe, each plugin should be independent.

if different plugins are using same function, we should push this function to public library

I found a better way to augument the request body.

membphis · 2024-09-25T03:25:37Z

apisix/plugins/ai-rag.lua

+
+local http     = require("resty.http")
+local core     = require("apisix.core")
+local decorate = require("apisix.plugins.ai-prompt-decorator").__decorate


membphis · 2024-09-25T03:26:49Z

apisix/plugins/ai-rag.lua

+local azure_openai_embeddings = require("apisix.plugins.ai-rag.embeddings.azure_openai").schema
+local azure_ai_search_schema = require("apisix.plugins.ai-rag.vector-search.azure_ai_search").schema
+
+local INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR


I suggest we keep the prefix HTTP_

membphis · 2024-09-25T03:27:33Z

apisix/plugins/ai-rag.lua

+
+local _M = {
+    version = 0.1,
+    priority = 1060, -- TODO check with other ai plugins


Do we need this comment TODO check with other ai plugins?

no we don't need this anymore

membphis · 2024-09-25T03:28:35Z

apisix/plugins/ai-rag/embeddings/azure_openai.lua

+-- limitations under the License.
+--
+local core = require("apisix.core")
+local INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR


ditto, keep prefix HTTP_

membphis · 2024-09-25T03:28:46Z

apisix/plugins/ai-rag/vector-search/azure_ai_search.lua

+-- limitations under the License.
+--
+local core = require("apisix.core")
+local INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR


membphis

I do not know much about this feature, only check the code style and how to use this plugin

membphis · 2024-09-25T08:36:37Z

docs/en/latest/plugins/ai-rag.md

+| **Field**                                       | **Required** | **Type** | **Description**                              |
+| ----------------------------------------------- | ------------ | -------- | -------------------------------------------- |
+| embeddings_provider                             | Yes          | Object   | Configuration for the embeddings provider    |
+| embeddings_provider.azure_openai                | Yes          | Object   | Configuration for Azure OpenAI embeddings    |


I have one question:

do we will support other provider?

yes, we might have to support other providers if the need be. So I have written code in a modular way.

Same issue, I think we'd better to add tip in doc for APISIX user, they should know this

docs/en/latest/plugins/ai-rag.md

kayx23 · 2024-09-29T06:23:22Z

docs/en/latest/plugins/ai-rag.md

+| vector_search_provider                          | Yes          | Object   | Configuration for the vector search provider |
+| vector_search_provider.azure_ai_search          | Yes          | Object   | Configuration for Azure AI Search            |
+| vector_search_provider.azure_ai_search.endpoint | No           | String   | Azure AI Search endpoint                     |
+| vector_search_provider.azure_ai_search.api_key  | No           | String   | Azure AI Search API key                      |


How come the Azure AI Search config details (endpoint and api key) are optional when their parent azure_ai_search is required?

oops, my bad.

docs/en/latest/plugins/ai-rag.md

kayx23 · 2024-09-29T07:05:39Z

docs/en/latest/plugins/ai-rag.md

+
+  | **Field** | **Required** | **Type** | **Description**              |
+  | --------- | ------------ | -------- | ---------------------------- |
+  | fields    | Yes          | String   | Fields for the vector search |


Casing consistency:

docs/en/latest/plugins/ai-rag.md

Co-authored-by: Traky Deng <trakydeng@gmail.com>

kayx23 · 2024-09-30T07:04:11Z

docs/en/latest/plugins/ai-rag.md

+
+The `ai-rag` plugin integrates Retrieval-Augmented Generation (RAG) capabilities with AI models.
+It allows efficient retrieval of relevant documents or information from external data sources and
+augments the AI's responses with that data, improving the accuracy and context of generated outputs.


Suggested change

augments the AI's responses with that data, improving the accuracy and context of generated outputs.

augments the LLM responses with that data, improving the accuracy and contextual relevance of the generated outputs.

By using the term LLM here you could probably save this entire sentence

kayx23 · 2024-09-30T07:05:12Z

docs/en/latest/plugins/ai-rag.md

+It allows efficient retrieval of relevant documents or information from external data sources and
+augments the AI's responses with that data, improving the accuracy and context of generated outputs.
+
+**_This plugin must be used in routes that proxy requests to LLMs only._**


same as #11541 (comment)

kayx23 · 2024-09-30T07:06:03Z

docs/en/latest/plugins/ai-rag.md

+
+**_This plugin must be used in routes that proxy requests to LLMs only._**
+
+**_As of now only Azure OpenAI and Azure AI Search services are supported for generating embeddings and performing vector search respectively, PRs for introducing support for other service providers are welcome._**


grammar fix and links

Suggested change

**_As of now only Azure OpenAI and Azure AI Search services are supported for generating embeddings and performing vector search respectively, PRs for introducing support for other service providers are welcome._**

**_As of now only [Azure OpenAI](https://azure.microsoft.com/en-us/products/ai-services/openai-service) and [Azure AI Search](https://azure.microsoft.com/en-us/products/ai-services/ai-search) services are supported for generating embeddings and performing vector search respectively. PRs for introducing support for other service providers are welcomed._**

kayx23 · 2024-09-30T07:06:50Z

docs/en/latest/plugins/ai-rag.md

+
+  | **Field** | **Required** | **Type** | **Description**              |
+  | --------- | ------------ | -------- | ---------------------------- |
+  | fields    | Yes          | String   | Fields for the vector search |


kayx23 · 2024-09-30T07:30:47Z

docs/en/latest/plugins/ai-rag.md

+        }
+      },
+      "override": {
+        "endpoint": "'"$azure_openai_endpoint"'"


the variables in the PUT request should be all caps as well to match the env vars above

zhoujiexiong · 2024-10-08T03:20:27Z

apisix/plugins/ai-rag/vector-search/azure_ai_search.lua

+        return nil, HTTP_INTERNAL_SERVER_ERROR, err
+    end
+
+    if res.status ~= 200 then


use ngx const var instead?

zhoujiexiong · 2024-10-08T03:21:28Z

apisix/plugins/ai-rag/embeddings/azure_openai.lua

+        return nil, HTTP_INTERNAL_SERVER_ERROR, err
+    end
+
+    if res.status ~= 200 then


use ngx const var instead?

zhoujiexiong · 2024-10-08T03:28:00Z

apisix/plugins/ai-rag.lua

+    local embeddings_provider_conf = conf.embeddings_provider[next(conf.embeddings_provider)]
+    local embeddings_driver = require("apisix.plugins.ai-rag.embeddings." .. embeddings_provider)
+
+    local vector_search_provider = next(conf.vector_search_provider)


also ref #11541 (comment) and adjust related?

zhoujiexiong · 2024-10-08T03:37:14Z

apisix/plugins/ai-rag.lua

+    local vs_req_schema = vector_search_driver.request_schema
+    local emb_req_schema = embeddings_driver.request_schema
+
+    request_schema.properties.ai_rag.properties.vector_search = vs_req_schema


I checked the process again and there is no risk of race condition. Ignore it.

bzp2010

There are no more comments, so if other reviewers agree with it, I'll agree with it.

docs/en/latest/plugins/ai-rag.md

zhoujiexiong

LGTM

shreemaan-abhishek commented Sep 11, 2024

View reviewed changes

apisix/plugins/ai-rag.lua Show resolved Hide resolved

shreemaan-abhishek commented Sep 11, 2024

View reviewed changes

apisix/plugins/ai-rag.lua Show resolved Hide resolved

shreemaan-abhishek commented Sep 11, 2024

View reviewed changes

initial commit

0c8e6b7

shreemaan-abhishek force-pushed the ai-rag branch from 6e5c04b to 0c8e6b7 Compare September 12, 2024 03:12

clean and test

e27ed8d

shreemaan-abhishek marked this pull request as ready for review September 13, 2024 08:29

dosubot bot added the size:XL This PR changes 500-999 lines, ignoring generated files. label Sep 13, 2024

clean

db4e663

dosubot bot added enhancement New feature or request plugin labels Sep 13, 2024

shreemaan-abhishek added 2 commits September 17, 2024 10:12

Merge branch 'master' of github.com:apache/apisix into ai-rag

8f22cdb

fix ci

4836c4d

zhoujiexiong suggested changes Sep 21, 2024

View reviewed changes

bzp2010 self-requested a review September 22, 2024 18:49

bzp2010 reviewed Sep 22, 2024

View reviewed changes

shreemaan-abhishek added 6 commits September 23, 2024 12:01

docs

e87c1e7

code review

ceafa66

code review

8e92bd3

f

ea2fcf6

update docs

bbe01a5

fix lint

edc3ccc

shreemaan-abhishek requested review from membphis and moonming September 25, 2024 02:53

membphis reviewed Sep 25, 2024

View reviewed changes

code review

d49665d

membphis reviewed Sep 25, 2024

View reviewed changes

add service provider related info

0e2bea2

kayx23 reviewed Sep 29, 2024

View reviewed changes

docs/en/latest/plugins/ai-rag.md Outdated Show resolved Hide resolved

kayx23 reviewed Sep 29, 2024

View reviewed changes

docs/en/latest/plugins/ai-rag.md Outdated Show resolved Hide resolved

kayx23 reviewed Sep 29, 2024

View reviewed changes

docs/en/latest/plugins/ai-rag.md Outdated Show resolved Hide resolved

kayx23 reviewed Sep 29, 2024

View reviewed changes

docs/en/latest/plugins/ai-rag.md Outdated Show resolved Hide resolved

kayx23 reviewed Sep 29, 2024

View reviewed changes

docs/en/latest/plugins/ai-rag.md Show resolved Hide resolved

kayx23 reviewed Sep 29, 2024

View reviewed changes

docs/en/latest/plugins/ai-rag.md Outdated Show resolved Hide resolved

shreemaan-abhishek and others added 5 commits September 30, 2024 11:12

Apply suggestions from code review

6761ffd

Co-authored-by: Traky Deng <trakydeng@gmail.com>

reformat

b17ce6f

uppercase

7701e05

make required in code too

458325c

Merge branch 'master' of github.com:apache/apisix into ai-rag

1d132fa

kayx23 reviewed Sep 30, 2024

View reviewed changes

docs: review comment

e16e873

zhoujiexiong suggested changes Oct 8, 2024

View reviewed changes

review comment

6f9d469

bzp2010 self-requested a review October 9, 2024 05:18

bzp2010 reviewed Oct 9, 2024

View reviewed changes

docs/en/latest/plugins/ai-rag.md Outdated Show resolved Hide resolved

shreemaan-abhishek added 2 commits October 9, 2024 12:24

remove vendor specific content

0f2f3ba

Merge branch 'master' of github.com:apache/apisix into ai-rag

fe3add2

zhoujiexiong approved these changes Oct 11, 2024

View reviewed changes

bzp2010 approved these changes Oct 12, 2024

View reviewed changes

shreemaan-abhishek requested review from kayx23 and membphis October 12, 2024 06:33

moonming approved these changes Oct 12, 2024

View reviewed changes

Revolyssup approved these changes Oct 16, 2024

View reviewed changes

shreemaan-abhishek merged commit 11c9d29 into apache:master Oct 16, 2024
48 checks passed


		request_schema.properties.ai_rag.properties.vector_search = vs_req_schema
		request_schema.properties.ai_rag.properties.embeddings = emb_req_schema

	if type(res_tab.data) ~= "table" or #res_tab.data < 1 then
	if type(res_tab.data) ~= "table" or core.table.isempty(res_tab.data) then

	local embeddings_provider_conf = conf.embeddings_provider[next(conf.embeddings_provider)]
	local embeddings_provider_conf = conf.embeddings_provider[embeddings_provider]

	local internal_server_error = ngx.HTTP_INTERNAL_SERVER_ERROR
	local INTERNAL_SERVER_ERROR = ngx.HTTP_INTERNAL_SERVER_ERROR

		@@ -114,4 +114,5 @@ function _M.rewrite(conf, ctx)
		end


		_M.__decorate = decorate -- for ai-rag plugin

	\| fields \| Yes \| String \| Fields for the vector search \|
	\| fields \| Yes \| string \| Fields for the vector search \|

	augments the AI's responses with that data, improving the accuracy and context of generated outputs.
	augments the LLM responses with that data, improving the accuracy and contextual relevance of the generated outputs.


		_This plugin must be used in routes that proxy requests to LLMs only._

		_As of now only Azure OpenAI and Azure AI Search services are supported for generating embeddings and performing vector search respectively, PRs for introducing support for other service providers are welcome._

	_As of now only Azure OpenAI and Azure AI Search services are supported for generating embeddings and performing vector search respectively, PRs for introducing support for other service providers are welcome._
	_As of now only [Azure OpenAI](https://azure.microsoft.com/en-us/products/ai-services/openai-service) and [Azure AI Search](https://azure.microsoft.com/en-us/products/ai-services/ai-search) services are supported for generating embeddings and performing vector search respectively. PRs for introducing support for other service providers are welcomed._

feat: ai-rag plugin #11568

feat: ai-rag plugin #11568

Conversation

shreemaan-abhishek commented Sep 11, 2024 • edited Loading

Description

Checklist

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zhoujiexiong left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bzp2010 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

membphis left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

kayx23 Sep 30, 2024 • edited Loading

Choose a reason for hiding this comment

kayx23 Sep 30, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bzp2010 left a comment

Choose a reason for hiding this comment

zhoujiexiong left a comment

Choose a reason for hiding this comment

shreemaan-abhishek commented Sep 11, 2024 •

edited

Loading

kayx23 Sep 30, 2024 •

edited

Loading

kayx23 Sep 30, 2024 •

edited

Loading