Update app.py

app.py

@@ -8,33 +8,42 @@ from report_generator import (
     generate_pdf_report,
     generate_diagram_report,
 )
-from transformers import pipeline
+import cohere
+from cohere.classify import ClassificationExample
 
-# Try to initialize the AI parser.
-try:
-    ai_parser = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")
-except Exception as e:
-    print("AI model failed to load, falling back to rule-based parsing:", e)
-    ai_parser = None
+# Initialize the Cohere client using your provided API key.
+# (For security, consider loading the API key from environment variables.)
+COHERE_API_KEY = "ffwtfyfLwqSOtw65psZardfqhUxMr1h7u5vzYotI"
+co = cohere.Client(COHERE_API_KEY)
 
-def ai_parse_command(command):
-    """Uses the AI model to classify the command intent; if unavailable, falls back to a heuristic."""
-    candidate_labels = ["report", "diagram"]
-    if ai_parser is not None:
-        try:
-            result = ai_parser(command, candidate_labels)
-            # Return the highest-scoring label.
-            return result["labels"][0]
-        except Exception as e:
-            print("AI parsing error, falling back to heuristic:", e)
-    # Fallback: use simple keyword search.
-    lower = command.lower()
-    if "diagram" in lower:
-        return "diagram"
-    elif "report" in lower:
-        return "report"
-    else:
-        return "unknown"
+# Define some labeled examples for classification.
+cohere_examples = [
+    ClassificationExample(text="generate a report on unemployment", label="report"),
+    ClassificationExample(text="create a report on market trends", label="report"),
+    ClassificationExample(text="generate diagram of sales data", label="diagram"),
+    ClassificationExample(text="create a diagram of quarterly revenue", label="diagram")
+]
+
+def cohere_parse_command(command):
+    """Uses Cohere's classify endpoint to determine the intent of the command."""
+    try:
+        response = co.classify(
+            model="large",
+            inputs=[command],
+            examples=cohere_examples
+        )
+        # Get the prediction for the first input.
+        prediction = response.classifications[0].prediction
+        return prediction
+    except Exception as e:
+        print("Cohere classification error, falling back to keyword search:", e)
+        lower = command.lower()
+        if "diagram" in lower:
+            return "diagram"
+        elif "report" in lower:
+            return "report"
+        else:
+            return "unknown"
 
 # Global containers for pending tasks and generated files.
 task_bucket = []      # Queue of tasks waiting to be processed.
@@ -48,8 +57,8 @@ def process_single_task(task):
     for remaining in range(5, 0, -1):
         task["timer"] = remaining
         time.sleep(1)
-    # Use AI (or fallback) to classify the intent.
-    intent = ai_parse_command(task["command"])
+    # Use Cohere (or fallback) to classify the intent.
+    intent = cohere_parse_command(task["command"])
     cmd = task["command"].lower()
 
     # Process based on the detected intent.