Update aws_aiclient.py

aws_aiclient.py

@@ -1,5 +1,4 @@
 # aws_aiclient.py
-
 import os
 import time
 import json
@@ -38,14 +37,18 @@ text_models = {
 
 class AIClient:
     def __init__(self):
-        self.client = ChatBedrockConverse(
-            model='meta.llama3-70b-instruct-v1:0',  # default model
+        self._model = 'meta.llama3-70b-instruct-v1:0'  # default model
+        self.client = self._create_client()
+        self.observability_manager = LLMObservabilityManager()
+        self.models = text_models
+
+    def _create_client(self, model: str = None) -> ChatBedrockConverse:
+        return ChatBedrockConverse(
+            model=model or self._model,
             region_name="ap-south-1",
             aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),
             aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY")
         )
-        self.observability_manager = LLMObservabilityManager()
-        self.models = text_models
 
     async def generate_response(
         self,
@@ -64,13 +67,13 @@ class AIClient:
         status = "success"
 
         try:
-            # Update the client's model if different from current
-            if model != self.client.model:
-                self.client.model = model
+            # Update the client if model is different
+            if model != self._model:
+                self._model = model
+                self.client = self._create_client()
 
             # Stream the response
             async for chunk in self.client.astream(messages):
-                print(chunk)
                 if chunk.content and chunk.content[0].get("text"):
                     content = chunk.content[0].get("text")
                     yield content