update sematic tool template

couchbaselabs · Feb 10, 2025 · 8b9347d · 8b9347d
1 parent 865b2d0
commit 8b9347d
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 12 deletions.
diff --git a/libs/agentc_core/agentc_core/tool/descriptor/secrets.py b/libs/agentc_core/agentc_core/tool/descriptor/secrets.py
@@ -13,8 +13,7 @@ class Couchbase(pydantic.BaseModel):
 
 
 class EmbeddingModelSecrets(pydantic.BaseModel):
-    class OpenAI(pydantic.BaseModel):
-        username: typing.Optional[str]
-        password: typing.Optional[str]
+    class EmbeddingModel(pydantic.BaseModel):
+        auth: str
 
-    openai: OpenAI
+    embedding: EmbeddingModel
diff --git a/libs/agentc_core/agentc_core/tool/generate/generator.py b/libs/agentc_core/agentc_core/tool/generate/generator.py
@@ -15,6 +15,8 @@
 from ..descriptor import HTTPRequestToolDescriptor
 from ..descriptor import SemanticSearchToolDescriptor
 from ..descriptor import SQLPPQueryToolDescriptor
+from ..descriptor.secrets import CouchbaseSecrets
+from ..descriptor.secrets import EmbeddingModelSecrets
 from .common import INPUT_MODEL_CLASS_NAME_IN_TEMPLATES
 from .common import OUTPUT_MODEL_CLASS_NAME_IN_TEMPLATES
 from .common import GeneratedCode
@@ -114,13 +116,26 @@ def generate(self) -> typing.Iterable[str]:
         with (self.template_directory / "semantic_q.jinja").open("r") as tmpl_fp:
             template = jinja2.Template(source=tmpl_fp.read(), autoescape=True)
             generation_time = datetime.datetime.now().strftime("%I:%M%p on %B %d, %Y")
+
+            cluster_secrets = None
+            for secret in self.record_descriptor.secrets:
+                if isinstance(secret, CouchbaseSecrets):
+                    cluster_secrets = secret
+                    break
+            embedding_secrets = None
+            for secret in self.record_descriptor.secrets:
+                if isinstance(secret, EmbeddingModelSecrets):
+                    embedding_secrets = secret
+                    break
+
             rendered_code = template.render(
                 {
                     "time": generation_time,
                     "tool": self.record_descriptor,
                     "input": input_model,
                     "vector_search": self.record_descriptor.vector_search,
-                    "secrets": self.record_descriptor.secrets[0].couchbase,
+                    "cluster_secrets": cluster_secrets.couchbase if cluster_secrets is not None else None,
+                    "embedding_secrets": embedding_secrets.embedding if embedding_secrets is not None else None,
                 }
             )
             logger.debug("The following code has been generated:\n" + rendered_code)

diff --git a/libs/agentc_core/agentc_core/tool/generate/templates/semantic_q.jinja b/libs/agentc_core/agentc_core/tool/generate/templates/semantic_q.jinja
@@ -21,11 +21,11 @@ logger = logging.getLogger(__name__)
 
 def _get_couchbase_cluster() -> couchbase.cluster.Cluster:
     authenticator = couchbase.auth.PasswordAuthenticator(
-        username=get_secret("{{ secrets.username }}").get_secret_value(),
-        password=get_secret("{{ secrets.password }}").get_secret_value(),
-        {% if secrets.certificate is none %}cert_path=get_secret("{{ secrets.certificate | safe }}").get_secret_value(){% endif %}
+        username=get_secret("{{ cluster_secrets.username }}").get_secret_value(),
+        password=get_secret("{{ cluster_secrets.password }}").get_secret_value(),
+        {% if cluster_secrets.certificate is none %}cert_path=get_secret("{{ cluster_secrets.certificate | safe }}").get_secret_value(){% endif %}
     )
-    conn_string = get_secret("{{ secrets.conn_string }}").get_secret_value()
+    conn_string = get_secret("{{ cluster_secrets.conn_string }}").get_secret_value()
     return couchbase.cluster.Cluster(conn_string, couchbase.options.ClusterOptions(authenticator))
 
 
@@ -51,17 +51,24 @@ def {{ tool.name }}(question: {{ input.type_name }}) -> list[str]:
     else:
         raise ValueError("Bad input given to tool!")
 
-    import sentence_transformers
-
     logger.debug("{{ tool.name }} is generating an embedding for: " + formatted_question + ".")
+    {% if vector_search.embedding_model.base_url is none %}
+    import sentence_transformers
     embedding_model = sentence_transformers.SentenceTransformer(
-        '{{ vector_search.embedding_model.name }}',
+        "{{ vector_search.embedding_model.name }}",
         tokenizer_kwargs={'clean_up_tokenization_spaces': True},
         cache_folder=DEFAULT_MODEL_CACHE_FOLDER,
         local_files_only=True
     )
     _embedding = embedding_model.encode(formatted_question)
     for_q = list(_embedding.astype('float64'))
+    {% else %}
+    import openai
+    open_ai_client = openai.OpenAI(
+        base_url="{{ vector_search.embedding_model.base_url }}", api_key=get_secret("{{ embedding_secrets.auth }}").get_secret_value()
+    )
+    for_q=open_ai_client.embeddings.create(model="{{ vector_search.embedding_model.name }}", input=formatted_question, encoding_format="float").data[0].embedding
+    {% endif %}
     vector_req = couchbase.vector_search.VectorSearch.from_vector_query(
         couchbase.vector_search.VectorQuery('{{ vector_search.vector_field }}', for_q, num_candidates={{ vector_search.num_candidates }})
     )