llumo 0.1.0__py3-none-any.whl → 0.1.3__py3-none-any.whl

llumo/.env

@@ -0,0 +1,6 @@
+
+BASE_URL="https://app.llumo.ai/api"
+postUrl = "https://red-skull-service-392377961931.us-central1.run.app/api/process-playground"
+fetchUrl = "https://red-skull-service-392377961931.us-central1.run.app/api/get-cells-data"
+validateUrl = "https://backend-api.llumo.ai/api/v1/workspace-key-details"
+SOCKET_URL="https://red-skull-service-392377961931.us-central1.run.app/"

llumo/__init__.py

@@ -1,2 +1,7 @@
 from .client import LlumoClient
 from .exceptions import LlumoAPIError
+from .helpingFuntions import *
+from .models import AVAILABLEMODELS
+from .execution import ModelExecutor
+
+

llumo/client.py

@@ -1,76 +1,551 @@
 import requests
 from .exceptions import LlumoAPIError
+import time
+import re
+import json
+import uuid
+import threading
+from .helpingFuntions import *
+from dotenv import load_dotenv
+import os
+import itertools
+import pandas as pd
+from typing import List, Dict
+from .models import AVAILABLEMODELS,getProviderFromModel
+from .execution import ModelExecutor
+from .sockets import LlumoSocketClient
 
-class LlumoClient:
-    """
-    A client to interact with Llumo API for evaluating AI-generated responses
-    """
+load_dotenv()  # Automatically looks for .env in current directory
+
+postUrl = os.getenv("postUrl")
+fetchUrl = os.getenv("fetchUrl")
+validateUrl = os.getenv("validateUrl")
+socketUrl = os.getenv("SOCKET_URL")
 
-    base_url = "https://app.llumo.ai/api"
+
+class LlumoClient:
 
     def __init__(self, api_key):
-        """
-        Initializes the LlumoClient with the given API key.
-
-        Parameters:
-        - api_key (str): The Llumo API key for authentication.
-        """
-        self.api_key = api_key
-
-
-    def EvaluateGrounded(self, outputText, groundTruthText, embeddingModelName="Google", metricsName="Cosine"):
-        """
-        Evaluates the groundedness of a response using a similarity metric.
-
-        Parameters:
-        - outputText (str): The generated output text to evaluate.
-        - groundTruthText (str): The reference ground truth text.
-        - embeddingModelName (str): Name of the embedding model to use. Default is "Google".
-        - metricsName (str): Similarity metric to apply (e.g., "Bleu"). Default is "Cosine".
-
-        Returns:
-        - dict: Contains statusCode, message, and evaluation data if successful.
-
-        Raises:
-        - LlumoAPIError for all specific error types.
-        """
-        url = f"{self.base_url}/external/grounded-external"
-
-        requestBody = {
-            "prompt": outputText,
-            "groundTruth": groundTruthText,
-            "embeddingModel": embeddingModelName,
-            "similarityMetric": metricsName,
-        }
+        self.apiKey = api_key
+        self.socket = LlumoSocketClient(socketUrl)
+        self.processMapping = {}
+
 
+    def validateApiKey(self, evalName = ""):
         headers = {
-            "Authorization": f"Bearer {self.api_key}",
-            "Content-Type": "application/json"
+            "Authorization": f"Bearer {self.apiKey}",
+            "Content-Type": "application/json",
         }
+        reqBody = {"analytics": [evalName]}
+
+        print(f"Making API key validation request to: {validateUrl}")
+        print(f"Request body: {reqBody}")
 
         try:
-            res = requests.post(url=url, json=requestBody, headers=headers)
+            response = requests.post(url=validateUrl, json=reqBody, headers=headers)
+            print(response.text)
+            # Print response info for debugging
+            print(f"Response status code: {response.status_code}")
+            print(f"Response headers: {response.headers}")
 
-            if res.status_code == 401:
-                raise LlumoAPIError.InvalidApiKey()
+            # Try to get at least some of the response content
+            try:
+                response_preview = response.text[:500]  # First 500 chars
+                print(f"Response preview: {response_preview}")
+            except Exception as e:
+                print(f"Could not get response preview: {e}")
 
-            res.raise_for_status()
-            result = res.json()
+        except requests.exceptions.RequestException as e:
+            print(f"Request exception: {str(e)}")
+            raise LlumoAPIError.RequestFailed(detail=str(e))
 
-            if 'data' not in result:
-                raise LlumoAPIError.InvalidApiResponse()
+        if response.status_code == 401:
+            raise LlumoAPIError.InvalidApiKey()
 
-            return {
-                "statusCode": result['data'].get('statusCode'),
-                "message": result['data'].get('message'),
-                "analytics": result['data']
-            }
+        # Handle other common status codes
+        if response.status_code == 404:
+            raise LlumoAPIError.RequestFailed(
+                detail=f"Endpoint not found (404): {validateUrl}"
+            )
 
-        except requests.exceptions.HTTPError as e:
-            raise LlumoAPIError.RequestFailed(str(e))
-        except ValueError:
+        if response.status_code >= 500:
+            raise LlumoAPIError.ServerError(
+                detail=f"Server error ({response.status_code})"
+            )
+
+        if response.status_code != 200:
+            raise LlumoAPIError.RequestFailed(
+                detail=f"Unexpected status code: {response.status_code}"
+            )
+
+        # Try to parse JSON
+        try:
+            data = response.json()
+        except ValueError as e:
+            print(f"JSON parsing error: {str(e)}")
+            print(
+                f"Response content that could not be parsed: {response.text[:1000]}..."
+            )
             raise LlumoAPIError.InvalidJsonResponse()
+
+        if "data" not in data or not data["data"]:
+            print(f"Invalid API response structure: {data}")
+            raise LlumoAPIError.InvalidApiResponse()
+
+        try:
+            self.hitsAvailable = data["data"].get("remainingHits", 0)
+            self.workspaceID = data["data"].get("workspaceID")
+            self.evalDefinition = data["data"].get("analyticsMapping")
+            self.token = data["data"].get("token")
+
+            print(f"API key validation successful:")
+            # print(f"- Remaining hits: {self.hitsAvailable}")
+            # print(f"- Workspace ID: {self.workspaceID}")
+            # print(f"- Token received: {'Yes' if self.token else 'No'}")
+
+        except Exception as e:
+            print(f"Error extracting data from response: {str(e)}")
+            raise LlumoAPIError.UnexpectedError(detail=str(e))
+
+    def postBatch(self, batch, workspaceID):
+        payload = {
+            "batch": json.dumps(batch),
+            "runType": "EVAL",
+            "workspaceID": workspaceID,
+        }
+        headers = {
+            "Authorization": f"Bearer {self.token}",
+            "Content-Type": "application/json",
+        }
+        try:
+            print(postUrl)
+            response = requests.post(postUrl, json=payload, headers=headers)
+            # print(f"Post API Status Code: {response.status_code}")
+            # print(response.text)
+
         except Exception as e:
-            raise LlumoAPIError.UnexpectedError(str(e))
+            print(f"Error in posting batch: {e}")
+
+    def AllProcessMapping(self):
+        for batch in self.allBatches:
+            for record in batch:
+                rowId = record['rowID']
+                colId = record['columnID']
+                pid = f'{rowId}-{colId}-{colId}'
+                self.processMapping[pid] = record
+
+
+    def finalResp(self,results):
+        seen = set()
+        uniqueResults = []
+
+        for item in results:
+            for rowID in item:  # Each item has only one key
+                if rowID not in seen:
+                    seen.add(rowID)
+                    uniqueResults.append(item)
+
+        return uniqueResults
+
+    def evaluate(self, dataframe, evals=["Response Completeness"],prompt_template = ""):
+        results = {}
+        try:
+            # Connect to socket first
+            print("Connecting to socket server...")
+            socketID = self.socket.connect(timeout=20)
+            print(f"Connected with socket ID: {socketID}")
+
+            # Process each evaluation
+            for eval in evals:
+                print(f"\n======= Running evaluation for: {eval} =======")
+
+                try:
+                    print(f"Validating API key for {eval}...")
+                    self.validateApiKey(evalName=eval)
+                    print(
+                        f"API key validation successful. Hits available: {self.hitsAvailable}"
+                    )
+                except Exception as e:
+                    print(f"Error during API key validation: {str(e)}")
+                    if (
+                        hasattr(e, "response")
+                        and getattr(e, "response", None) is not None
+                    ):
+                        print(f"Status code: {e.response.status_code}")
+                        print(f"Response content: {e.response.text[:500]}...")
+                    raise
+
+                if self.hitsAvailable == 0 or len(dataframe) > self.hitsAvailable:
+                    raise LlumoAPIError.InsufficientCredits()
+
+                evalDefinition = self.evalDefinition[eval]
+                model = "GPT_4"
+                provider = "OPENAI"
+                evalType = "LLM"
+                workspaceID = self.workspaceID
+
+                # Prepare all batches before sending
+                print("Preparing batches...")
+                self.allBatches = []
+                currentBatch = []
+
+                for index, row in dataframe.iterrows():
+
+                    tools = row["tools"] if "tools" in dataframe.columns else []
+                    groundTruth = row["groundTruth"] if "groundTruth" in dataframe.columns else ""
+                    messageHistory = row["messageHistory"] if "messageHistory" in dataframe.columns else []
+                    promptTemplate = prompt_template
+
+                    keys = re.findall(r"{{(.*?)}}", promptTemplate)
+
+                    # extracting the required values for the the columns based on the prompt template
+                    inputDict = {key: row[key] for key in keys if key in row}
+                    output = row["output"]
+
+                    activePlayground = (
+                        f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+                    )
+                    rowID = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+                    columnID = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace(
+                        "-", ""
+                    )
+
+                    # Use the server-provided socket ID here
+                    templateData = {
+                        "processID": getProcessID(),
+                        "socketID": socketID,  # Using the server-assigned socket ID
+                        "processData": {
+                            "executionDependency": {
+                                "query": "",
+                                "context": "",
+                                "output": output,
+                                "tools": tools,
+                                "groundTruth": groundTruth,
+                                "messageHistory": messageHistory,
+                            },
+                            "definition": evalDefinition,
+                            "model": model,
+                            "provider": provider,
+                            "analytics": eval,
+                        },
+                        "workspaceID": workspaceID,
+                        "type": "EVAL",
+                        "evalType": evalType,
+                        "kpi": eval,
+                        "columnID": columnID,
+                        "rowID": rowID,
+                        "playgroundID": activePlayground,
+                        "processType": "EVAL",
+                    }
+
+                    # Build query/context from input
+                    query = ""
+                    context = ""
+                    for key, value in inputDict.items():
+                        if isinstance(value, str):
+                            length = len(value.split()) * 1.5
+                            if length > 50:
+                                context += f" {key}: {value}, "
+                            else:
+                                if promptTemplate:
+                                    tempObj = {key: value}
+                                    promptTemplate = getInputPopulatedPrompt(promptTemplate, tempObj)
+                                else:
+                                    query += f" {key}: {value}, "
+
+                    if not context.strip():
+                        for key, value in inputDict.items():
+                            context += f" {key}: {value}, "
+
+                    templateData["processData"]["executionDependency"]["context"] = context.strip()
+                    templateData["processData"]["executionDependency"]["query"] = query.strip()
+
+                    if promptTemplate and not query.strip():
+                        templateData["processData"]["executionDependency"]["query"] = promptTemplate
+
+                    currentBatch.append(templateData)
+
+                    if len(currentBatch) == 10 or index == len(dataframe) - 1:
+                        self.allBatches.append(currentBatch)
+                        currentBatch = []
+
+                # Post all batches
+                total_items = sum(len(batch) for batch in self.allBatches)
+                print(f"Posting {len(self.allBatches)} batches ({total_items} items total)")
+
+                for cnt, batch in enumerate(self.allBatches):
+                    print(
+                        f"Posting batch {cnt + 1}/{len(self.allBatches)} for eval '{eval}'"
+                    )
+                    try:
+                        self.postBatch(batch=batch, workspaceID=workspaceID)
+                        print(f"Batch {cnt + 1} posted successfully")
+                    except Exception as e:
+                        print(f"Error posting batch {cnt + 1}: {str(e)}")
+                        continue
+
+
+
+                    # Small delay between batches to prevent overwhelming the server
+                    time.sleep(1)
+
+                # updating the dict for row column mapping
+                self.AllProcessMapping()
+                # Calculate a reasonable timeout based on the data size
+                timeout = max(60, min(600, total_items * 10))
+                print(
+                    f"All batches posted. Waiting up to {timeout} seconds for results..."
+                )
+
+                # Listen for results
+                self.socket.listen_for_results(
+                    min_wait=10, max_wait=timeout, inactivity_timeout=30
+                )
+
+                # Get results for this evaluation
+                eval_results = self.socket.get_received_data()
+                print(f"Received {len(eval_results)} results for evaluation '{eval}'")
+
+                # Add these results to our overall results
+                results[eval] = self.finalResp(eval_results)
+                print(f"======= Completed evaluation: {eval} =======\n")
+
+            print("All evaluations completed successfully")
+
+        except Exception as e:
+            print(f"Error during evaluation: {e}")
+            raise
+        finally:
+            # Always disconnect the socket when done
+            try:
+                self.socket.disconnect()
+                print("Socket disconnected")
+            except Exception as e:
+                print(f"Error disconnecting socket: {e}")
+
+        for evalName, records in results.items():
+            for item in records:
+                self.processMapping[list(item.keys())[0]] = list(item.values())[0]
+
+
+
+            dataframe[evalName] =  self.processMapping.values()
+
+        return dataframe
+
+    def evaluateCompressor(self, dataframe, prompt_template):
+        results = []
+        try:
+            # Connect to socket first
+            print("Connecting to socket server...")
+            socketID = self.socket.connect(timeout=20)
+            print(f"Connected with socket ID: {socketID}")
+
+            try:
+                print(f"Validating API key...")
+                self.validateApiKey()
+                print(f"API key validation successful. Hits available: {self.hitsAvailable}")
+            except Exception as e:
+                print(f"Error during API key validation: {str(e)}")
+                if hasattr(e, "response") and getattr(e, "response", None) is not None:
+                    print(f"Status code: {e.response.status_code}")
+                    print(f"Response content: {e.response.text[:500]}...")
+                raise
+
+            if self.hitsAvailable == 0 or len(dataframe) > self.hitsAvailable:
+                raise LlumoAPIError.InsufficientCredits()
+
+            model = "GPT_4"
+            provider = "OPENAI"
+            evalType = "LLUMO"
+            workspaceID = self.workspaceID
+
+            # Prepare all batches before sending
+            print("Preparing batches...")
+            self.allBatches = []
+            currentBatch = []
+
+            for index, row in dataframe.iterrows():
+                promptTemplate = prompt_template
+
+                # extracting the placeholders from the prompt template
+                keys = re.findall(r"{{(.*?)}}", promptTemplate)
+                inputDict = {key: row[key] for key in keys if key in row}
+
+                activePlayground = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+                rowID = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+                columnID = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+
+                compressed_prompt_id = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+                compressed_prompt_output_id = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+                cost_id = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+                cost_saving_id = f"{int(time.time() * 1000)}{uuid.uuid4()}".replace("-", "")
+
+                # Use the server-provided socket ID here
+                templateData = {
+                    "processID": getProcessID(),
+                    "socketID": socketID,
+                    "rowID": rowID,
+                    "columnID": columnID,
+                    "processType": "COST_SAVING",
+                    "evalType": evalType,
+                    "dependency": list(inputDict.keys()),
+                    "costColumnMapping": {
+                        "compressed_prompt": compressed_prompt_id,
+                        "compressed_prompt_output": compressed_prompt_output_id,
+                        "cost": cost_id,
+                        "cost_saving": cost_saving_id
+                    },
+                    "processData": {
+                        "rowData": {
+                            "query": {"type": "VARIABLE", "value": ""},
+                            "context": {"type": "VARIABLE", "value": ""},
+                        },
+                        "dependency": list(inputDict.keys()),
+                        "dependencyMapping": {ky: ky for ky in list(inputDict.keys())},
+                        "provider": provider,
+                        "model": model,
+                        "promptText": promptTemplate,
+                        "costColumnMapping": {
+                            "compressed_prompt": compressed_prompt_id,
+                            "compressed_prompt_output": compressed_prompt_output_id,
+                            "cost": cost_id,
+                            "cost_saving": cost_saving_id
+                        }
+                    },
+                    "workspaceID": workspaceID,
+                    "email": "",
+                    "playgroundID": activePlayground
+                }
+
+
+                # Build query/context from input
+                query = ""
+                context = ""
+
+                for key, value in inputDict.items():
+                    if isinstance(value, str):
+                        length = len(value.split()) * 1.5
+                        if length > 50:
+                            context += f" {key}: {value}, "
+                        else:
+                            if promptTemplate:
+                                populatedPrompt = getInputPopulatedPrompt(promptTemplate, {key: value})
+                                query += f"{populatedPrompt} "
+                            else:
+                                query += f" {key}: {value}, "
+
+                if not context.strip():
+                    for key, value in inputDict.items():
+                        context += f" {key}: {value}, "
+
+                templateData["processData"]["rowData"]["context"]["value"] = context.strip()
+                templateData["processData"]["rowData"]["query"]["value"] = query.strip()
+
+                if promptTemplate and not query.strip():
+                    templateData["processData"]["rowData"]["query"]["value"] = promptTemplate
+
+                print(templateData)
+                currentBatch.append(templateData)
+
+            if len(currentBatch) == 10 or index == len(dataframe) - 1:
+                    self.allBatches.append(currentBatch)
+                    currentBatch = []
+
+            # Post all batches
+            total_items = sum(len(batch) for batch in self.allBatches)
+            print(f"Posting {len(self.allBatches)} batches ({total_items} items total)")
+
+            for cnt, batch in enumerate(self.allBatches):
+                print(f"Posting batch {cnt + 1}/{len(self.allBatches)} for eval '{eval}'")
+                try:
+                    self.postBatch(batch=batch, workspaceID=workspaceID)
+                    print(f"Batch {cnt + 1} posted successfully")
+                except Exception as e:
+                    print(f"Error posting batch {cnt + 1}: {str(e)}")
+                    continue
+
+                # Small delay between batches to prevent overwhelming the server
+                time.sleep(1)
+
+            # updating the dict for row column mapping
+            self.AllProcessMapping()
+            # Calculate a reasonable timeout based on the data size
+            timeout = max(60, min(600, total_items * 10))
+            print(f"All batches posted. Waiting up to {timeout} seconds for results...")
+
+            # Listen for results
+            self.socket.listen_for_results(min_wait=10, max_wait=timeout, inactivity_timeout=30)
+
+            # Get results for this evaluation
+            eval_results = self.socket.get_received_data()
+            print(f"Received {len(eval_results)} results for evaluation '{eval}'")
+
+            # Add these results to our overall results
+            results = self.finalResp(eval_results)
+            print(f"======= Completed evaluation: {eval} =======\n")
+
+            print("All evaluations completed successfully")
+
+        except Exception as e:
+            print(f"Error during evaluation: {e}")
+            raise
+        finally:
+            # Always disconnect the socket when done
+            try:
+                self.socket.disconnect()
+                print("Socket disconnected")
+            except Exception as e:
+                print(f"Error disconnecting socket: {e}")
+
+        compressed_prompt , compressed_prompt_output  , cost , cost_saving  = costColumnMapping(results,self.processMapping)
+        dataframe["compressed_prompt"] = compressed_prompt
+        dataframe["compressed_prompt_output"] = compressed_prompt_output
+        dataframe["cost"] = cost
+        dataframe["cost_saving"] = cost_saving
+        return dataframe
+
+
+    def run_sweep(self,templates: List[str], dataset: Dict[str, List[str]], model_aliases: List[AVAILABLEMODELS], apiKey: str, evals = ["Response Correctness"]) -> pd.DataFrame:
+        executor = ModelExecutor(apiKey)
+
+        keys = list(dataset.keys())
+        value_combinations = list(itertools.product(*dataset.values()))
+        combinations = [dict(zip(keys, values)) for values in value_combinations]
+
+        results = []
 
+        # Iterate through combinations
+        for combo in combinations:
+            for template in templates:
+                prompt = template
+                for k, v in combo.items():
+                    prompt = prompt.replace(f"{{{{{k}}}}}", v)
+                # Add a row for each model
+                for model in model_aliases:
+                    row = {
+                        "template": template,
+                        "prompt": prompt,
+                        **combo,
+                        "model": model.value
+                    }
 
+                    
+                    try:
+                        provider = getProviderFromModel(model)
+                        response = executor.execute(provider, model.value, prompt, apiKey)
+                        row["output"] = response
+                    except Exception as e:
+                        row["output"] = f"Error: {str(e)}"
+                    
+                    results.append(row)
+        df=pd.DataFrame(results)
+        df.to_csv("sweep_results.csv", index=False)
+        print(str(templates[0]))
+        res = self.evaluate(df,evals =evals,prompt_template=str(templates[0]))
+        return res
+    
+    
+class SafeDict(dict):
+    def __missing__(self, key):
+        return ""

llumo/exceptions.py

@@ -28,3 +28,4 @@ class LlumoAPIError(Exception):
     @staticmethod
     def EvalError(detail="Some error occured while processing"):
         return LlumoAPIError(f"error: {detail}")
+

llumo/execution.py

@@ -0,0 +1,39 @@
+import openai
+import google.generativeai as genai
+from .models import Provider
+
+class ModelExecutor:
+    def __init__(self, apiKey: str):
+        self.apiKey = apiKey
+
+    def execute(self, provider: Provider, modelName: str, prompt: str,api_key) -> str:
+        if provider == Provider.OPENAI:
+            return self._executeOpenAI(modelName, prompt,api_key)
+        elif provider == Provider.GOOGLE:
+            return self._executeGoogle(modelName, prompt,api_key)
+        else:
+            raise ValueError(f"Unsupported provider: {provider}")
+
+    def _executeOpenAI(self, modelName: str, prompt: str,api_key) -> str:
+        client = openai.OpenAI(api_key=api_key)
+        response = client.chat.completions.create(model="gpt-4",  # Replace with the desired model
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": prompt}  # User's prompt
+            ]
+        )
+        return response.choices[0].message.content
+
+    def _executeGoogle(self, modelName: str, prompt: str,api_key) -> str:
+        
+        # Configure GenAI with API Key
+        genai.configure(api_key=api_key)
+            
+        # Select Generative Model
+        model = genai.GenerativeModel("gemini-1.5-flash-latest")
+        # Generate Response
+        response = model.generate_content(prompt)
+        return response.text
+        
+    
+ 

llumo/helpingFuntions.py

@@ -0,0 +1,60 @@
+import time
+import uuid
+
+def getProcessID():
+    return f"{int(time.time() * 1000)}{uuid.uuid4()}"
+
+
+def getInputPopulatedPrompt(promptTemplate, tempObj):
+    for key, value in tempObj.items():
+        promptTemplate = promptTemplate.replace(f"{{{{{key}}}}}", value)
+    return promptTemplate
+
+
+
+import time
+import uuid
+
+def getProcessID():
+    return f"{int(time.time() * 1000)}{uuid.uuid4()}"
+
+
+def getInputPopulatedPrompt(promptTemplate, tempObj):
+    for key, value in tempObj.items():
+        promptTemplate = promptTemplate.replace(f"{{{{{key}}}}}", value)
+    return promptTemplate
+
+def costColumnMapping(costResults,allProcess):
+    # this dict will store cost column data for each row
+    cost_cols = {}
+    compressed_prompt = []
+    compressed_prompt_output = []
+    cost = []
+    cost_saving = []
+    print("BATCHES: ",allProcess)
+    print("COST RESULTS :", costResults)
+    # iterate through each batch
+    for record in allProcess:
+        cost_cols[record] = []
+        # iterate through each record of cost saving results received from the api
+        for item in costResults:
+            # fetching all cost column data for a specific row. i.e each row will have 4 columns
+            if list(item.keys())[0].split("-")[0] == record.split("-")[0]:
+                cost_cols[record].append(list(item.values())[0])
+
+    for ky, val in cost_cols.items():
+        # compressed prompt column
+        compressed_prompt.append(val[0])
+        # compressed output
+        compressed_prompt_output.append(val[1])
+        # cost
+        cost.append(val[2])
+        # cost saved
+        cost_saving.append(val[3])
+
+    return compressed_prompt , compressed_prompt_output  , cost , cost_saving
+
+
+
+
+

llumo/models.py

@@ -0,0 +1,43 @@
+from enum import Enum
+
+class Provider(str, Enum):
+    OPENAI = "OPENAI"
+    GOOGLE = "GOOGLE"
+
+# Maps model aliases → (provider, actual model name for API)
+_MODEL_METADATA = {
+    "GPT_4": (Provider.OPENAI, "gpt-4"),
+    "GPT_4_32K": (Provider.OPENAI, "gpt-4-32k"),
+    "GPT_35T": (Provider.OPENAI, "gpt-3.5-turbo"),
+    "GPT_35T_INS": (Provider.OPENAI, "gpt-3.5-turbo-instruct"),
+    "GPT_35T_16K": (Provider.OPENAI, "gpt-3.5-turbo-16k"),
+    "GPT_35_TURBO": (Provider.OPENAI, "gpt-3.5-turbo"),
+
+    "GOOGLE_15_FLASH": (Provider.GOOGLE, "gemini-1.5-flash-latest"),
+    "GEMINI_PRO": (Provider.GOOGLE, "gemini-pro"),
+    "TEXT_BISON": (Provider.GOOGLE, "text-bison-001"),
+    "CHAT_BISON": (Provider.GOOGLE, "chat-bison-001"),
+    "TEXT_BISON_32K": (Provider.GOOGLE, "text-bison-32k"),
+    "TEXT_UNICORN": (Provider.GOOGLE, "text-unicorn-experimental"),
+}
+
+class AVAILABLEMODELS(str, Enum):
+    GPT_4 = "gpt-4"
+    GPT_4_32K = "gpt-4-32k"
+    GPT_35T = "gpt-3.5-turbo"
+    GPT_35T_INS = "gpt-3.5-turbo-instruct"
+    GPT_35T_16K = "gpt-3.5-turbo-16k"
+    GPT_35_TURBO = "gpt-3.5-turbo"
+
+    GOOGLE_15_FLASH = "gemini-1.5-flash-latest"
+    GEMINI_PRO = ""
+    TEXT_BISON = "text-bison-001"
+    CHAT_BISON = "chat-bison-001"
+    TEXT_BISON_32K = "text-bison-32k"
+    TEXT_UNICORN = "text-unicorn-experimental"
+
+def getProviderFromModel(model: AVAILABLEMODELS) -> Provider:
+    for alias, (provider, apiName) in _MODEL_METADATA.items():
+        if model.value == apiName:
+            return provider
+    raise ValueError(f"Provider not found for model: {model}")

llumo/sockets.py

@@ -0,0 +1,154 @@
+import socketio
+import threading
+import time
+
+
+class LlumoSocketClient:
+    def __init__(self, socket_url):
+        self.socket_url = socket_url
+        self._received_data = []
+        self._last_update_time = None
+        self._listening_done = threading.Event()
+        self._connection_established = threading.Event()
+        self._lock = threading.Lock()
+        self._connected = False
+        self.server_socket_id = None  # Store the server-assigned socket ID
+
+        # Initialize client
+        self.sio = socketio.Client(
+            # logger=True,
+            # engineio_logger=True,
+            reconnection=True,
+            reconnection_attempts=5,
+            reconnection_delay=1,
+        )
+
+        @self.sio.on("connect")
+        def on_connect():
+            print("Socket connection established")
+            self._connected = True
+            # Don't set connection_established yet - wait for server confirmation
+
+        # Listen for the connection-established event from the server
+        @self.sio.on("connection-established")
+        def on_connection_established(data):
+            print(
+                f"Server acknowledged connection with 'connection-established' event: {data}"
+            )
+            if isinstance(data, dict) and "socketId" in data:
+                self.server_socket_id = data["socketId"]
+                print(f"Received server socket ID: {self.server_socket_id}")
+            self._connection_established.set()
+
+        @self.sio.on("result-update")
+        def on_result_update(data):
+            with self._lock:
+                print(f"Received result-update event: {data}")
+                self._received_data.append(data)
+                self._last_update_time = time.time()
+
+        @self.sio.on("disconnect")
+        def on_disconnect():
+            print("Socket disconnected")
+            self._connected = False
+
+        @self.sio.on("connect_error")
+        def on_connect_error(error):
+            print(f"Socket connection error: {error}")
+
+        @self.sio.on("error")
+        def on_error(error):
+            print(f"Socket error event: {error}")
+
+    def connect(self, timeout=20):
+        self._received_data = []
+        self._connection_established.clear()
+        self._listening_done.clear()
+        self.server_socket_id = None
+
+        try:
+            print("Attempting direct WebSocket connection...")
+            # Connect with websocket transport
+            self.sio.connect(self.socket_url, transports=["websocket"], wait=True)
+
+            print(f"Engine.IO connection established with SID: {self.sio.sid}")
+            print(
+                "Waiting for server to acknowledge connection with connection-established event..."
+            )
+
+            # Wait for the connection-established event
+            if not self._connection_established.wait(timeout):
+                raise RuntimeError("Timed out waiting for connection-established event")
+
+            self._last_update_time = time.time()
+            print(
+                f"Connection fully established. Server socket ID: {self.server_socket_id}"
+            )
+
+            # Return the server-assigned socket ID if available, otherwise fall back to the client's SID
+            return self.server_socket_id or self.sio.sid
+        except Exception as e:
+            self._connected = False
+            raise RuntimeError(f"WebSocket connection failed: {e}")
+
+    def listen_for_results(self, min_wait=5, max_wait=300, inactivity_timeout=30):
+        """
+        Listen for results with improved timeout handling:
+        - min_wait: Minimum time to wait even if no data is received
+        - max_wait: Maximum total time to wait for results
+        - inactivity_timeout: Time to wait after last data received
+        """
+        if not self._connected:
+            raise RuntimeError("WebSocket is not connected. Call connect() first.")
+
+        start_time = time.time()
+        self._last_update_time = time.time()
+
+        def timeout_watcher():
+            while not self._listening_done.is_set():
+                current_time = time.time()
+                time_since_last_update = current_time - self._last_update_time
+                total_elapsed = current_time - start_time
+
+                # Always wait for minimum time
+                if total_elapsed < min_wait:
+                    time.sleep(0.5)
+                    continue
+
+                # Stop if maximum time exceeded
+                if total_elapsed > max_wait:
+                    print(
+                        f"⚠️ Maximum wait time of {max_wait}s reached, stopping listener."
+                    )
+                    self._listening_done.set()
+                    break
+
+                # Stop if no activity for inactivity_timeout
+                if time_since_last_update > inactivity_timeout:
+                    print(
+                        f"⚠️ No data received for {inactivity_timeout}s, stopping listener."
+                    )
+                    self._listening_done.set()
+                    break
+
+                # Check every second
+                time.sleep(1)
+
+        timeout_thread = threading.Thread(target=timeout_watcher, daemon=True)
+        timeout_thread.start()
+        print("Started listening for WebSocket events...")
+        self._listening_done.wait()
+        print(f"Finished listening. Received {len(self._received_data)} data updates.")
+
+    def get_received_data(self):
+        with self._lock:
+            return self._received_data.copy()
+
+    def disconnect(self):
+        try:
+            if self._connected:
+                self.sio.disconnect()
+                self._connected = False
+                print("WebSocket client disconnected")
+        except Exception as e:
+            print(f"Error during WebSocket disconnect: {e}")

{llumo-0.1.0.dist-info → llumo-0.1.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: llumo
-Version: 0.1.0
+Version: 0.1.3
 Summary: Python SDK for interacting with the Llumo ai API.
 Home-page: https://www.llumo.ai/
 Author: Llumo
@@ -8,11 +8,14 @@ Author-email: product@llumo.ai
 License: Proprietary
 Requires-Python: >=3.7
 License-File: LICENSE
-Requires-Dist: requests>=2.25.1
-Requires-Dist: setuptools>=58.1.0
-Requires-Dist: twine>=6.1.0
-Requires-Dist: wheel>=0.45.1
-Requires-Dist: build>=1.2.2.post1
+Requires-Dist: requests>=2.0.0
+Requires-Dist: websocket-client>=1.0.0
+Requires-Dist: pandas>=1.0.0
+Requires-Dist: numpy>=1.0.0
+Requires-Dist: python-socketio[client]==5.13.0
+Requires-Dist: python-dotenv==1.1.0
+Requires-Dist: openai==1.75.0
+Requires-Dist: google-generativeai==0.8.5
 Dynamic: author
 Dynamic: author-email
 Dynamic: home-page

llumo-0.1.3.dist-info/RECORD

@@ -0,0 +1,13 @@
+llumo/.env,sha256=Vx5FkuywpYHXH2N8epJ7PlNOPiwx9UP9DUz4vWd0urs,373
+llumo/__init__.py,sha256=8ZgAtxJNNgHorEXoxaLQ2YWrVXGgamoayyLMD1L4FbE,183
+llumo/client.py,sha256=DggiOLmBG21lEpg1vqjV5SC-PfR2LuVnpsY6HMTyF9I,23086
+llumo/exceptions.py,sha256=l3_5d9cBMm-hwpuFrg3nvI9cEP2GTKXcCyWiWHwnYDM,1041
+llumo/execution.py,sha256=ZvbZDSAvwj1XwSlgPNiy4r9fZG_vtfSlaWGwNI9xCa8,1453
+llumo/helpingFuntions.py,sha256=HPy2w3IaYfH_hDBgXdoAmNZmAbDUO01bgW7gHBGNw8A,1765
+llumo/models.py,sha256=WBtnu7ckOy9TGRiwswz04xOGYF6EslTUOxHUz4QWzUA,1602
+llumo/sockets.py,sha256=M6piy6bNt342GmTQCdUJJDUgMYGxk0Acjgj11uI4Vdg,5965
+llumo-0.1.3.dist-info/licenses/LICENSE,sha256=vMiqSi3KpDHq3RFxKiqdh10ZUF3PjE3nnntANU-HEu4,186
+llumo-0.1.3.dist-info/METADATA,sha256=DyqkkQAIg95hiKyFdYrHj_CazvTb8ocSUZLi13cslLc,721
+llumo-0.1.3.dist-info/WHEEL,sha256=pxyMxgL8-pra_rKaQ4drOZAegBVuX-G_4nRHjjgWbmo,91
+llumo-0.1.3.dist-info/top_level.txt,sha256=d5zUTMI99llPtLRB8rtSrqELm_bOqX-bNC5IcwlDk88,6
+llumo-0.1.3.dist-info/RECORD,,

{llumo-0.1.0.dist-info → llumo-0.1.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (79.0.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

llumo-0.1.0.dist-info/RECORD

@@ -1,8 +0,0 @@
-llumo/__init__.py,sha256=WT0Y6g-MtJqIIMpH1AoWQfL-XraaQDtMbjApjBHRPqk,72
-llumo/client.py,sha256=w1rd6uT4F-iwHty8xDqXy3Q5BM_9u1jgdSZvy032LQE,2468
-llumo/exceptions.py,sha256=BQcLqfViMxPklGIPJnH1tfajdytvuGpd5Sidv4ta6h0,1039
-llumo-0.1.0.dist-info/licenses/LICENSE,sha256=vMiqSi3KpDHq3RFxKiqdh10ZUF3PjE3nnntANU-HEu4,186
-llumo-0.1.0.dist-info/METADATA,sha256=OcJtsjap9U3XwzxjR03V4VzITQbPkOO_JgCcSgg0qo0,593
-llumo-0.1.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-llumo-0.1.0.dist-info/top_level.txt,sha256=d5zUTMI99llPtLRB8rtSrqELm_bOqX-bNC5IcwlDk88,6
-llumo-0.1.0.dist-info/RECORD,,