Update llama.py

llama.py

@@ -8,18 +8,10 @@ import logging
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
-# Load environment variables
+# Load environment variables (only needed if using a local .env file)
 load_dotenv()
-API_KEYS = [
-    os.getenv("GORQ_API_KEY_1"),
-    os.getenv("GORQ_API_KEY_2"),
-    os.getenv("GORQ_API_KEY_3"),
-    os.getenv("GORQ_API_KEY_4")
-]
 
 class GorQClient:
-    """Client for interacting with the Groq API to call the LLaMA model."""
-
     def __init__(self):
         self.current_api_index = 0
         self.client = Groq(api_key=self.get_api_key())
@@ -27,17 +19,21 @@ class GorQClient:
 
     def switch_api_key(self):
         """Switch to the next available API key."""
-        self.current_api_index = (self.current_api_index + 1) % len(API_KEYS)
+        self.current_api_index = (self.current_api_index + 1) % self.number_of_api_keys()
         logging.info("Switched to API key index: %d", self.current_api_index)
         self.client = Groq(api_key=self.get_api_key())
 
     def get_api_key(self):
-        """Get the current API key."""
-        return API_KEYS[self.current_api_index]
+        """Get the current API key from environment variables."""
+        return os.getenv(f"GORQ_API_KEY_{self.current_api_index + 1}")
+
+    def number_of_api_keys(self):
+        """Get the number of available API keys."""
+        return 4  # Change this if you add more keys
 
     def call_llama_model(self, user_message):
         """Call the LLaMA model and return the response."""
-        max_retries = len(API_KEYS)
+        max_retries = self.number_of_api_keys()
         retries = 0
         
         while retries < max_retries:
@@ -71,12 +67,14 @@ class GorQClient:
                 for chunk in completion:
                     response += chunk.choices[0].delta.content or ""
 
+                # Clean up the response to remove formatting artifacts
                 response = self.clean_response(response)
+
                 logging.info("Response received from LLaMA model.")
                 return response
 
             except Exception as e:
-                logging.error(f"API error with key index {self.current_api_index}: {e}")
+                logging.error(f"API error: {e}")
                 self.switch_api_key()
                 retries += 1
                 logging.info("Retrying with new API key.")
@@ -85,10 +83,8 @@ class GorQClient:
 
     def clean_response(self, response):
         """Remove any unwanted characters or formatting artifacts."""
-        unwanted_chars = ["*", "•", "#", "`"]
-        for char in unwanted_chars:
-            response = response.replace(char, "")
-        return response.strip()
+        response = response.replace("*", "").replace("•", "").replace("#", "").replace("`", "").strip()
+        return response
 
 def format_confidence_level(confidence):
     """Convert numerical confidence to descriptive text."""
@@ -106,6 +102,7 @@ def generate_response_based_on_yolo(yolo_output):
     if not yolo_output:
         return "Fracture Analysis Report: No fractures detected in the provided image. However, if you're experiencing pain or discomfort, please consult a healthcare professional for a thorough evaluation."
 
+    # Process YOLO output to create detailed analysis
     current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
     details = []
     
@@ -130,6 +127,7 @@ def generate_response_based_on_yolo(yolo_output):
             }
         })
 
+    # Construct detailed message for LLM
     user_message = json.dumps({
         "timestamp": current_time,
         "analysis_type": "Fracture Detection and Analysis",
@@ -137,10 +135,12 @@ def generate_response_based_on_yolo(yolo_output):
         "request": "Provide a comprehensive analysis including immediate care recommendations, pain management strategies, and follow-up care instructions. Include specific details about each detected fracture and potential complications to watch for."
     }, indent=2)
 
+    # Get detailed response from LLM
     response = gorq_client.call_llama_model(user_message)
     
     return response
 
+# Test the functionality (This part should be executed in a separate test file or interactive environment)
 if __name__ == "__main__":
     # Example YOLO output for testing
     example_yolo_output = [
@@ -152,5 +152,6 @@ if __name__ == "__main__":
         }
     ]
 
+    # Generate a response based on the example YOLO output
     response = generate_response_based_on_yolo(example_yolo_output)
     print(response)