Update llm_handler.py

llm_handler.py

@@ -1,9 +1,13 @@
 from openai import OpenAI
 from params import OPENAI_MODEL, OPENAI_API_KEY
+import llamanet
 
 # Create an instance of the OpenAI class
 client = OpenAI(api_key=OPENAI_API_KEY)
 
+# Initialize LlamaNet client
+llamanet_client = llamanet.Client()
+
 def send_to_chatgpt(msg_list):
     try:
         completion = client.chat.completions.create(
@@ -25,8 +29,34 @@ def send_to_chatgpt(msg_list):
         print(f"Error in send_to_chatgpt: {str(e)}")
         return f"Error: {str(e)}", None
 
+def send_to_llamanet(msg_list):
+    try:
+        # Convert msg_list to the format expected by LlamaNet
+        llamanet_messages = [{"role": msg["role"], "content": msg["content"]} for msg in msg_list]
+        
+        # Send request to LlamaNet
+        response = llamanet_client.chat.completions.create(
+            model="llamanet",
+            messages=llamanet_messages,
+            stream=True
+        )
+        
+        llamanet_response = ""
+        for chunk in response:
+            if chunk.choices[0].delta.content is not None:
+                llamanet_response += chunk.choices[0].delta.content
+        
+        # LlamaNet doesn't provide usage information
+        llamanet_usage = None
+        return llamanet_response, llamanet_usage
+    except Exception as e:
+        print(f"Error in send_to_llamanet: {str(e)}")
+        return f"Error: {str(e)}", None
+
 def send_to_llm(provider, msg_list):
     if provider == "llamanet":
+        return send_to_llamanet(msg_list)
+    elif provider == "openai":
         return send_to_chatgpt(msg_list)
     else:
         raise ValueError(f"Unknown provider: {provider}")