Couchbase Async

1import asyncio
2import os
3import time
4
5from kern.agent import Agent
6from kern.knowledge.embedder.openai import OpenAIEmbedder
7from kern.knowledge.knowledge import Knowledge
8from kern.vectordb.couchbase import CouchbaseSearch
9from couchbase.auth import PasswordAuthenticator
10from couchbase.management.search import SearchIndex
11from couchbase.options import ClusterOptions, KnownConfigProfiles
12
13# Couchbase connection settings
14username = os.getenv("COUCHBASE_USER")  # Replace with your username
15password = os.getenv("COUCHBASE_PASSWORD")  # Replace with your password
16connection_string = os.getenv("COUCHBASE_CONNECTION_STRING")
17
18# Create cluster options with authentication
19auth = PasswordAuthenticator(username, password)
20cluster_options = ClusterOptions(auth)
21cluster_options.apply_profile(KnownConfigProfiles.WanDevelopment)
22
23# Define the vector search index
24search_index = SearchIndex(
25    name="vector_search",
26    source_type="gocbcore",
27    idx_type="fulltext-index",
28    source_name="recipe_bucket",
29    plan_params={"index_partitions": 1, "num_replicas": 0},
30    params={
31        "doc_config": {
32            "docid_prefix_delim": "",
33            "docid_regexp": "",
34            "mode": "scope.collection.type_field",
35            "type_field": "type",
36        },
37        "mapping": {
38            "default_analyzer": "standard",
39            "default_datetime_parser": "dateTimeOptional",
40            "index_dynamic": True,
41            "store_dynamic": True,
42            "default_mapping": {"dynamic": True, "enabled": False},
43            "types": {
44                "recipe_scope.recipes": {
45                    "dynamic": False,
46                    "enabled": True,
47                    "properties": {
48                        "content": {
49                            "enabled": True,
50                            "fields": [
51                                {
52                                    "docvalues": True,
53                                    "include_in_all": False,
54                                    "include_term_vectors": False,
55                                    "index": True,
56                                    "name": "content",
57                                    "store": True,
58                                    "type": "text",
59                                }
60                            ],
61                        },
62                        "embedding": {
63                            "enabled": True,
64                            "dynamic": False,
65                            "fields": [
66                                {
67                                    "vector_index_optimized_for": "recall",
68                                    "docvalues": True,
69                                    "dims": 3072,
70                                    "include_in_all": False,
71                                    "include_term_vectors": False,
72                                    "index": True,
73                                    "name": "embedding",
74                                    "similarity": "dot_product",
75                                    "store": True,
76                                    "type": "vector",
77                                }
78                            ],
79                        },
80                        "meta": {
81                            "dynamic": True,
82                            "enabled": True,
83                            "properties": {
84                                "name": {
85                                    "enabled": True,
86                                    "fields": [
87                                        {
88                                            "docvalues": True,
89                                            "include_in_all": False,
90                                            "include_term_vectors": False,
91                                            "index": True,
92                                            "name": "name",
93                                            "store": True,
94                                            "analyzer": "keyword",
95                                            "type": "text",
96                                        }
97                                    ],
98                                }
99                            },
100                        },
101                    },
102                }
103            },
104        },
105    },
106)
107
108knowledge_base = Knowledge(
109    vector_db=CouchbaseSearch(
110        bucket_name="recipe_bucket",
111        scope_name="recipe_scope",
112        collection_name="recipes",
113        couchbase_connection_string=connection_string,
114        cluster_options=cluster_options,
115        search_index=search_index,
116        embedder=OpenAIEmbedder(
117            id="text-embedding-3-large",
118            dimensions=3072,
119            api_key=os.getenv("OPENAI_API_KEY"),
120        ),
121        wait_until_index_ready=60,
122        overwrite=True,
123    ),
124)
125
126# Create and use the agent
127agent = Agent(knowledge=knowledge_base)
128
129async def run_agent():
130    await knowledge_base.ainsert(
131        url="https://kern-public.s3.amazonaws.com/recipes/ThaiRecipes.pdf"
132    )
133    time.sleep(5)  # wait for the vector index to be sync with kv
134    await agent.aprint_response("How to make Thai curry?", markdown=True)
135
136if __name__ == "__main__":
137    asyncio.run(run_agent())