Update app.py

app.py

@@ -111,6 +111,27 @@ def generate_synthetic_data(description, columns):
         return generated_text
     except Exception as e:
         return f"Error: {e}"
+        
+def clean_generated_text(generated_text):
+    # Extract CSV part using a regular expression
+    csv_match = re.search(r'(\n?([A-Za-z0-9_]+,)*[A-Za-z0-9_]+\n([^\n,]*,)*[^\n,]*\n*)+', generated_text)
+    
+    if csv_match:
+        csv_text = csv_match.group(0)
+    else:
+        raise ValueError("No valid CSV data found in generated text.")
+    
+    return csv_text
+
+def process_generated_data(csv_data):
+    # Clean the generated data
+    cleaned_data = clean_generated_text(csv_data)
+    
+    # Convert to DataFrame
+    data = StringIO(cleaned_data)
+    df = pd.read_csv(data)
+    
+    return df
 
 @app.post("/generate/")
 def generate_data(request: DataGenerationRequest):
@@ -125,11 +146,7 @@ def generate_data(request: DataGenerationRequest):
     df_synthetic = process_generated_data(generated_data)
     return JSONResponse(content={"data": df_synthetic.to_dict(orient="records")})
 
-def process_generated_data(csv_data):
-    csv_data = csv_data.strip()
-    data = StringIO(csv_data)
-    df = pd.read_csv(data)
-    return df
+
 
 @app.get("/")
 def greet_json():