KIG / kig_core /utils.py
heymenn's picture
Upload 12 files
ec6d5f9 verified
import pandas as pd
import io
import logging
from typing import List
from .schemas import KeyIssue # Import the Pydantic model
logger = logging.getLogger(__name__)
def key_issues_to_dataframe(key_issues: List[KeyIssue]) -> pd.DataFrame:
"""Converts a list of KeyIssue objects into a Pandas DataFrame."""
if not key_issues:
return pd.DataFrame()
# Use Pydantic's .model_dump() for robust serialization
data = [ki.model_dump() for ki in key_issues]
df = pd.DataFrame(data)
# Optional: Reorder or rename columns if needed
# df = df[['id', 'title', 'description', 'challenges', 'potential_impact']] # Example reordering
return df
def dataframe_to_excel_bytes(df: pd.DataFrame) -> bytes:
"""Converts a Pandas DataFrame to Excel format in memory (bytes)."""
logger.info("Generating Excel file from DataFrame...")
output = io.BytesIO()
try:
# Use BytesIO object as the target file
with pd.ExcelWriter(output, engine='openpyxl') as writer:
df.to_excel(writer, index=False, sheet_name='Key Issues')
excel_data = output.getvalue()
logger.info("Excel file generated successfully.")
return excel_data
except Exception as e:
logger.error(f"Failed to generate Excel file: {e}", exc_info=True)
raise RuntimeError("Failed to create Excel output.") from e
# Removed: format_df (HTML specific, less relevant for Excel output)
# Removed: init_app (handled by config.py)
# Removed: get_model (handled by llm_interface.py)
# Removed: clear_memory (handle state/memory management within LangGraph setup if needed)
# Removed: _set_env (handled by config.py and dotenv)
# Kept: format_doc (renamed to format_doc_for_llm in graph_operations.py)
# Removed: update_doc_history (reducer logic should be handled in LangGraph state definition/nodes)