From e5578cd79093cfa1b99c2d15170afb9e5f1db0c2 Mon Sep 17 00:00:00 2001
From: Matthew Keeler <mkeeler@launchdarkly.com>
Date: Wed, 27 Nov 2024 09:00:08 -0500
Subject: [PATCH] fix: Fix context usage for message interpolation

---
 ldai/client.py                    |  2 +-
 ldai/testing/test_model_config.py | 41 ++++++++++++++++++++++++++++---
 2 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/ldai/client.py b/ldai/client.py
index 1eceeac..6a802dd 100644
--- a/ldai/client.py
+++ b/ldai/client.py
@@ -152,7 +152,7 @@ def config(
         all_variables = {}
         if variables:
             all_variables.update(variables)
-        all_variables['ldctx'] = context
+        all_variables['ldctx'] = context.to_dict()
 
         messages = None
         if 'messages' in variation and isinstance(variation['messages'], list) and all(
diff --git a/ldai/testing/test_model_config.py b/ldai/testing/test_model_config.py
index 593a037..afab1e4 100644
--- a/ldai/testing/test_model_config.py
+++ b/ldai/testing/test_model_config.py
@@ -43,7 +43,19 @@ def td() -> TestData:
         .variations(
             {
                 'model': {'id': 'fakeModel', 'parameters': {'extra-attribute': 'I can be anything I set my mind/type to'}},
-                'messages': [{'role': 'system', 'content': 'Hello, {{ldctx.name}}!'}],
+                'messages': [{'role': 'system', 'content': 'Hello, {{ldctx.name}}! Is your last name {{ldctx.last}}?'}],
+                '_ldMeta': {'enabled': True, 'versionKey': 'abcd'},
+            }
+        )
+        .variation_for_all(0)
+    )
+
+    td.update(
+        td.flag('multi-ctx-interpolation')
+        .variations(
+            {
+                'model': {'id': 'fakeModel', 'parameters': {'extra-attribute': 'I can be anything I set my mind/type to'}},
+                'messages': [{'role': 'system', 'content': 'Hello, {{ldctx.user.name}}! Do you work for {{ldctx.org.shortname}}?'}],
                 '_ldMeta': {'enabled': True, 'versionKey': 'abcd'},
             }
         )
@@ -192,7 +204,7 @@ def test_provider_config_handling(ldai_client: LDAIClient):
 
 
 def test_context_interpolation(ldai_client: LDAIClient):
-    context = Context.builder('user-key').name("Sandy").build()
+    context = Context.builder('user-key').name("Sandy").set('last', 'Beaches').build()
     default_value = AIConfig(enabled=True, model=ModelConfig('fake-model'), messages=[])
     variables = {'name': 'World'}
 
@@ -202,7 +214,30 @@ def test_context_interpolation(ldai_client: LDAIClient):
 
     assert config.messages is not None
     assert len(config.messages) > 0
-    assert config.messages[0].content == 'Hello, Sandy!'
+    assert config.messages[0].content == 'Hello, Sandy! Is your last name Beaches?'
+    assert config.enabled is True
+
+    assert config.model is not None
+    assert config.model.id == 'fakeModel'
+    assert config.model.get_parameter('temperature') is None
+    assert config.model.get_parameter('maxTokens') is None
+    assert config.model.get_parameter('extra-attribute') == 'I can be anything I set my mind/type to'
+
+
+def test_multi_context_interpolation(ldai_client: LDAIClient):
+    user_context = Context.builder('user-key').name("Sandy").build()
+    org_context = Context.builder('org-key').kind('org').name("LaunchDarkly").set('shortname', 'LD').build()
+    context = Context.multi_builder().add(user_context).add(org_context).build()
+    default_value = AIConfig(enabled=True, model=ModelConfig('fake-model'), messages=[])
+    variables = {'name': 'World'}
+
+    config, _ = ldai_client.config(
+        'multi-ctx-interpolation', context, default_value, variables
+    )
+
+    assert config.messages is not None
+    assert len(config.messages) > 0
+    assert config.messages[0].content == 'Hello, Sandy! Do you work for LD?'
     assert config.enabled is True
 
     assert config.model is not None