feat: knowledge base tool (do not review now)

[yuye-aws]

Description

POC for knowledge base tool. You can use the knowledge base tool with the following steps:

# update connector settings
PUT _cluster/settings
{
  "persistent": {
    "plugins.ml_commons.trusted_connector_endpoints_regex": [
      """^https://runtime\.sagemaker\..*[a-z0-9-]\.amazonaws\.com/.*$""",
      """^https://api\.openai\.com/.*$""",
      """^https://api\.cohere\.ai/.*$""",
      """^https://bedrock-runtime\..*[a-z0-9-]\.amazonaws\.com/.*$""",
      """^https://bedrock-agent-runtime\..*[a-z0-9-]\.amazonaws\.com/.*$"""
    ]
  }
}

# register bedrock knowledge base connector
POST /_plugins/_ml/connectors/_create
{
  "name": "Amazon Bedrock Connector: knowledge",
  "description": "The connector to the Bedrock knowledge base",
  "version": 1,
  "protocol": "aws_sigv4",
  "parameters": {
    "region": "...",
    "service_name": "bedrock"
  },
  "credential": {
    "access_key": "...",
    "secret_key": "..."
  },
  "actions": [
    {
      "action_type": "predict",
      "method": "POST",
      "url": "...",
      "headers": {
        "content-type": "application/json",
        "x-amz-content-sha256": "required"
      },
      "request_body": "{\"retrievalQuery\": {\"text\": \"${parameters.text}\"}}"
    }
  ]
}

# register remote model given knowledge base connector
POST /_plugins/_ml/models/_register?deploy=true
{
  "name": "bedrock: knowledge base",
  "function_name": "remote",
  "description": "Test connector for bedrock knowledge base",
  "connector_id": "..."
}

# test your remote knowledge base model
POST /_plugins/_ml/models/<model_id>/_predict/
{
  "parameters": {
    "text": "..."
  }
}

# register a flow agent with knowledge base tool
POST /_plugins/_ml/agents/_register
{
  "name": "Test_Agent_For_KB_TOOL",
  "type": "flow",
  "description": "this is a test agent for knowledge base tool",
  "tools": [
    {
      "type": "KnowledgeBaseTool",
      "parameters": {
        "model_id": "..."
      }
    }
  ]
} 

# test your knowledge base agent
POST /_plugins/_ml/agents/<agent_id>/_execute
{
  "parameters": {
    "text": "..."
  }
}

Related Issues

POC for knowledge base tool

Check List

• New functionality includes testing.
• New functionality has been documented.
• API changes companion pull request created.
• Commits are signed per the DCO using --signoff.
• Public documentation issue/PR created.

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
For more information on following Developer Certificate of Origin and signing off your commits, please check here.

[xinyual]

Can you please add UT and IT?

[xinyual]

If knowledge base endpoint returns an error like exceed limitation (I don't know, we may need a survey), we will return empty result, does this meet expectation?

[yuye-aws]

Any idea to mock an error exception from bedrock knowledge base?

[conggguan]

Any idea to mock an error exception from bedrock knowledge base?

We can make a easy real environment exception by remove the model access to the embedding model and then call the API.
Maybe return the origin exception information is enough?

[yuye-aws]

Good idea. I can start testing once we have plan to merge this PR. I will mark this PR as draft.

[yuye-aws]

Can you please add UT and IT?

It's a POC for now. We are not in a hurry to merge this PR, so I will ad UT and IT later.

[conggguan]

How should we assure that the input parameter model_id is a valid knowledge base connector?

[yuye-aws]

How should we assure that the input parameter model_id is a valid knowledge base connector?

A simple idea is: fetch model and connector configuration from the model_id. Then, check the url from the connector. Just wondering whether PPLTool or other tools has this kind of model_id validation @xinyual

[Hailong-am]

can we leverage existing MLModelTool or ConnectorTool

[codecov]

Codecov Report

Attention: Patch coverage is 0% with 62 lines in your changes missing coverage. Please review.

Project coverage is 77.02%. Comparing base (2b76e3c) to head (4033781).
Report is 6 commits behind head on main.

Files	Patch %	Lines
.../org/opensearch/agent/tools/KnowledgeBaseTool.java	0.00%	59 Missing ⚠️
src/main/java/org/opensearch/agent/ToolPlugin.java	0.00%	3 Missing ⚠️

Additional details and impacted files

@@             Coverage Diff              @@
##               main     #353      +/-   ##
============================================
- Coverage     81.78%   77.02%   -4.77%     
- Complexity      193      222      +29     
============================================
  Files            11       16       +5     
  Lines           961     1197     +236     
  Branches        137      171      +34     
============================================
+ Hits            786      922     +136     
- Misses          121      211      +90     
- Partials         54       64      +10     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[yuye-aws]

No need to review this PR now. This PR is just a POC for alert analysis workflow.