使用 Kong AI 网关和 Redis 进行语义处理和向量相似性搜索

from sentence_transformers import SentenceTransformer
from sentence_transformers.util import truncate_embeddings

model = SentenceTransformer('all-mpnet-base-v2', cache_folder="./")
embeddings = model.encode("Who is Joseph Conrad?")
embeddings = truncate_embeddings(embeddings, 3)
print(embeddings.size)
print(embeddings)

3
[ 0.06030013 -0.00782523  0.01018228]

from sentence_transformers import SentenceTransformer

model = SentenceTransformer('all-mpnet-base-v2', cache_folder="./")

sentences = [
    "Who is Joseph Conrad?",
    "Tell me more about Joseph Conrad.",
    "Living is easy with eyes closed.",
]

embeddings = model.encode(sentences)
print(embeddings.shape)

similarities = model.similarity(embeddings, embeddings)
print(similarities)

(3, 768)
tensor([[1.0000, 0.8600, 0.0628],
        [0.8600, 1.0000, 0.1377],
        [0.0628, 0.1377, 1.0000]])

import redis
from redis.commands.search.field import TextField, VectorField
from redis.commands.search.indexDefinition import IndexDefinition, IndexType
from redis.commands.search.query import Query
import numpy as np
import openai
import os

### Get environment variables
openai.api_key = os.getenv("OPENAI_API_KEY")
host = os.getenv("REDIS_LB")

### Create a Redis Index for the Vector Embeddings
client = redis.Redis(host=host, port=6379)

try:
    client.ft('index1').dropindex(delete_documents=True)
except:
    print("index does not exist")


schema = (
    TextField("name"),
    TextField("description"),
    VectorField(
        "vector",
        "FLAT",
        {
            "TYPE": "FLOAT32",
            "DIM": 1536,
            "DISTANCE_METRIC": "COSINE",
        }
    ),
)

definition = IndexDefinition(prefix=["vectors:"], index_type=IndexType.HASH)
res = client.ft("index1").create_index(fields=schema, definition=definition)


### Step 1: call OpenAI to generate Embeddings for the reference text and stores it in Redis

name = "vector1"
content = "Who is Joseph Conrad?"
redis_key = f"vectors:{name}"

res = openai.embeddings.create(input=content, model="text-embedding-3-small").data[0].embedding

embeddings = np.array(res, dtype=np.float32).tobytes()

pipe = client.pipeline()
pipe.hset(redis_key, mapping = {
  "name": name,
  "description": content,
  "vector": embeddings
})
res = pipe.execute()


### Step 2: perform Vector Range queries with 2 new texts and get the distance (similarity) score

query = (
    Query("@vector:[VECTOR_RANGE $radius $vec]=>{$yield_distance_as: distance_score}")
     .return_fields("id", "distance_score")
     .dialect(2)
)

# Text #1
content = "Tell me more about Joseph Conrad"
res = openai.embeddings.create(input=content, model="text-embedding-3-small").data[0].embedding
new_embeddings = np.array(res, dtype=np.float32).tobytes()

query_params = {
    "radius": 1,
    "vec": new_embeddings
}
res = client.ft("index1").search(query, query_params).docs
print(res)

# Text #2
content = "Living is easy with eyes closed"
res = openai.embeddings.create(input=content, model="text-embedding-3-small").data[0].embedding
new_embeddings = np.array(res, dtype=np.float32).tobytes()

query_params = {
    "radius": 1,
    "vec": new_embeddings
}
res = client.ft("index1").search(query, query_params).docs
print(res)

While executing the code, you can monitor Redis with, for example, redis-cli monitor. The code line res = client.ft("index1").search(query, query_params).docs 

should log a message like this one:

"FT.SEARCH" "index1" "@vector:[VECTOR_RANGE $radius $vec]=>{$YIELD_DISTANCE_AS: score}" "RETURN" "2" "id" "score" "DIALECT" "2" 
"LIMIT" "0" "10" "params" "4" "radius" "1" "vec" "\xcb9\x9c<\xf8T\x18=\xaa\xd4\xb5\xbcB\xc0.=\xb5………."

query = (
    Query("@vector:[VECTOR_RANGE $radius $vec]=>{$yield_distance_as: distance_score}")
     .return_fields("id", "distance_score")
     .dialect(2)
)

query_params = {
    "radius": 1,
    "vec": new_embeddings
}

[Document {'id': 'vectors:vector1', 'payload': None, 'distance_score': '0.123970687389'}]
[Document {'id': 'vectors:vector1', 'payload': None, 'distance_score': '0.903066933155'}]

_format_version: "3.0"
_info:
  select_tags:
  - semantic-cache
_konnect:
  control_plane_name: default
services:
- name: service1
  host: localhost
  port: 32000
  routes:
  - name: route1
    paths:
    - /openai-route
    plugins:
    - name: ai-proxy
      instance_name: ai-proxy-openai-route
      enabled: true
      config:
        auth:
          header_name: Authorization
          header_value: Bearer <your_OPENAI_APIKEY>
        route_type: llm/v1/chat
        model:
          provider: openai
          name: gpt-4
          options:
            max_tokens: 512
            temperature: 1.0
    - name: ai-semantic-cache
      instance_name: ai-semantic-cache-openai
      enabled: true
      config:
        embeddings:
          auth:
            header_name: Authorization
            header_value: Bearer <your_OPENAI_APIKEY>
          model:
            provider: openai
            name: text-embedding-3-small
            options:
              upstream_url: https://api.openai.com/v1/embeddings
        vectordb:
          dimensions: 1536
          distance_metric: cosine
          strategy: redis
          threshold: 0.2
          redis:
            host: redis-stack.redis.svc.cluster.local
            port: 6379

curl -i -X POST \
  --url $DATA_PLANE_LB/openai-route \
  --header 'Content-Type: application/json' \
   --data '{
   "messages": [
     {
       "role": "user",
       "content": "Who is Joseph Conrad?"
     }
   ]
 }'

HTTP/1.1 200 OK
Content-Type: application/json
Connection: keep-alive
X-Cache-Status: Miss
x-ratelimit-limit-requests: 10000
CF-RAY: 8fce86cde915eae2-ORD
x-ratelimit-limit-tokens: 10000
x-ratelimit-remaining-requests: 9999
x-ratelimit-remaining-tokens: 9481
x-ratelimit-reset-requests: 8.64s
x-ratelimit-reset-tokens: 3.114s
access-control-expose-headers: X-Request-ID
x-request-id: req_29afd8838136a2f7793d6c129430b341
X-Content-Type-Options: nosniff
openai-organization: user-4qzstwunaw6d1dhwnga5bc5q
Date: Sat, 04 Jan 2025 22:05:00 GMT
alt-svc: h3=":443"; ma=86400
openai-processing-ms: 10002
openai-version: 2020-10-01
CF-Cache-Status: DYNAMIC
strict-transport-security: max-age=31536000; includeSubDomains; preload
Server: cloudflare
Content-Length: 1456
X-Kong-LLM-Model: openai/gpt-4
X-Kong-Upstream-Latency: 10097
X-Kong-Proxy-Latency: 471
Via: 1.1 kong/3.9.0.0-enterprise-edition
X-Kong-Request-Id: 36f6b41df3b74f78f586ae327af27075

{
  "id": "chatcmpl-Am6YEtvUquPHdHdcI59eZC3UfOUVz",
  "object": "chat.completion",
  "created": 1736028290,
  "model": "gpt-4-0613",
  "choices": [
    {
      "index": 0,
      "message": {
        "role": "assistant",
        "content": "Joseph Conrad was a Polish-British writer regarded as one of the greatest novelists to write in the English language. He was born on December 3, 1857, and died on August 3, 1924. Though he did not speak English fluently until his twenties, he was a master prose stylist who brought a non-English sensibility into English literature.\n\nConrad wrote stories and novels, many with a nautical setting, that depict trials of the human spirit in the midst of what he saw as an impassive, inscrutable universe. His notable works include \"Heart of Darkness\", \"Lord Jim\", and \"Nostromo\". Conrad's writing often presents a deep, pessimistic view of the world and deals with the theme of the clash of cultures and moral ambiguity.",
        "refusal": null
      },
      "logprobs": null,
      "finish_reason": "stop"
    }
  ],
  "usage": {
    "prompt_tokens": 12,
    "completion_tokens": 163,
    "total_tokens": 175,
    "prompt_tokens_details": {
      "cached_tokens": 0,
      "audio_tokens": 0
    },
    "completion_tokens_details": {
      "reasoning_tokens": 0,
      "audio_tokens": 0,
      "accepted_prediction_tokens": 0,
      "rejected_prediction_tokens": 0
    }
  },
  "system_fingerprint": null
}

"FT.CREATE" "idx:vss_kong_semantic_cache:511efd84-117b-4c89-87cb-f92f9b74a6c0:openai-gpt-4" "ON" "JSON" "PREFIX" "1" "kong_semantic_cache:511efd84-117b-4c89-87cb-f92f9b74a6c0:openai-gpt-4:" "SCORE" "1.0" "SCHEMA" "$.vector" "AS" "vector" "VECTOR" "FLAT" "6" "TYPE" "FLOAT32" "DIM" "1536" "DISTANCE_METRIC" "COSINE"

"FT.SEARCH""idx:vss_kong_semantic_cache:511efd84-117b-4c89-87cb-f92f9b74a6c0:openai-gpt-4" "@vector:[VECTOR_RANGE $range $query_vector]=>{$YIELD_DISTANCE_AS: vector_score}" "SORTBY" "vector_score" "DIALECT" "2" "LIMIT" "0" "4" "PARAMS" "4" "query_vector" "\x981\x87<bE\xe4<b\xa3\..........\xbc" "range" "0.2"

curl -i -X POST \
  --url $DATA_PLANE_LB/openai-route \
  --header 'Content-Type: application/json' \
   --data '{
   "messages": [
     {
       "role": "user",
       "content": "Tell me more about Joseph Conrad"
     }
   ]
 }'
HTTP/1.1 200 OK
Date: Sun, 05 Jan 2025 14:28:59 GMT
Content-Type: application/json; charset=utf-8
Connection: keep-alive
X-Cache-Status: Hit
Age: 0
X-Cache-Key: kong_semantic_cache:511efd84-117b-4c89-87cb-f92f9b74a6c0:openai-gpt-4:fcdf7d8995a227392f839b4530f8d8c3055748b96275fa9558523619172fd2a8
X-Cache-Ttl: 288
Content-Length: 1020
X-Kong-Response-Latency: 221
Server: kong/3.9.0.0-enterprise-edition
X-Kong-Request-Id: eef1373a3a688a68f088a52f72318315

{"object":"chat.completion","system_fingerprint":null,"id":"fcdf7d8995a22…….

curl -i -X POST \
  --url $DATA_PLANE_LB/openai-route \
  --header 'Content-Type: application/json' \
   --data '{
   "messages": [
     {
       "role": "user",
       "content": "Living is easy with eyes closed"
     }
   ]
 }'

redis-cli --scan
"kong_semantic_cache:511efd84-117b-4c89-87cb-f92f9b74a6c0:openai-gpt-4:fcdf7d8995a227392f839b4530f8d8c3055748b96275fa9558523619172fd2a8"
"kong_semantic_cache:511efd84-117b-4c89-87cb-f92f9b74a6c0:openai-gpt-4:22fbee1a1a45147167f29cc53183d0d2eef618c973e4284ad0179970209cf131"

_format_version: "3.0"
_info:
  select_tags:
  - semantic-routing
_konnect:
  control_plane_name: default
services:
- name: service1
  host: localhost
  port: 32000
  routes:
  - name: route1
    paths:
    - /openai-route
    plugins:
    - name: ai-proxy-advanced
      instance_name: ai-proxy-openai-route
      enabled: true
      config:
        balancer:
          algorithm: semantic
        embeddings:
          auth:
            header_name: Authorization
            header_value: Bearer <your_OPENAI_APIKEY>
          model:
            provider: openai
            name: text-embedding-3-small
            options:
              upstream_url: "https://api.openai.com/v1/embeddings"
        vectordb:
          dimensions: 1536
          distance_metric: cosine
          strategy: redis
          threshold: 0.8
          redis:
            host: redis-stack.redis.svc.cluster.local
            port: 6379
        targets:
        - model:
            provider: openai
            name: gpt-4
          route_type: "llm/v1/chat"
          auth:
            header_name: Authorization
            header_value: Bearer <your_OPENAI_APIKEY>
          description: "mathematics, algebra, calculus, trigonometry"
        - model:
            provider: openai
            name: gpt-4o-mini
          route_type: "llm/v1/chat"
          auth:
            header_name: Authorization
            header_value: Bearer <your_OPENAI_APIKEY>
          description: "piano, orchestra, liszt, classical music"

% curl -s -X POST \
  --url $DATA_PLANE_LB/openai-route \
  --header 'Content-Type: application/json' \
  --data '{
    "messages": [ 
      { 
        "role": "user", 
        "content": "Who wrote the Hungarian Rhapsodies piano pieces?"
      } 
    ] 
  }' | jq '.model' 
"gpt-4o-mini-2024-07-18"

% curl -s -X POST \
  --url $DATA_PLANE_LB/openai-route \
  --header 'Content-Type: application/json' \
  --data '{
    "messages": [
      {
        "role": "user",
        "content": "Tell me a contemporary pianist of Chopin"
      }
    ]
  }' | jq '.model'
"gpt-4o-mini-2024-07-18"

"gpt-4-0613".

% curl -s -X POST \
  --url $DATA_PLANE_LB/openai-route \
  --header 'Content-Type: application/json' \
  --data '{
     "messages": [
       {
         "role": "user",
         "content": "Tell me about Fermat''s last theorem"
       }
     ]
   }' | jq '.model'
"gpt-4-0613"