app.py

import gradio as gr
import pandas as pd
import numpy as np
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import FAISS
from langchain.prompts import ChatPromptTemplate
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
import os
from dotenv import load_dotenv

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")

# Load CSV data
df = pd.read_csv("K-Funding250331.csv")

# Clean and prepare data
def clean_and_prepare_data(df):
    # Fill NaN values with empty strings
    df_clean = df.fillna("")
    
    # Create document-like entries for each row
    documents = []
    for idx, row in df_clean.iterrows():
        # Convert the row to a string representation
        doc_content = f"Startup: {row['Startup Name']}\n"
        for col in df_clean.columns:
            if col != 'Startup Name' and row[col]:  # Add only non-empty fields
                doc_content += f"{col}: {row[col]}\n"
        
        documents.append({"content": doc_content, "source": f"row_{idx}"})
    
    return documents

# Convert DataFrame to documents
documents = clean_and_prepare_data(df)

# Split documents into chunks
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=100
)
splits = text_splitter.create_documents([doc["content"] for doc in documents])

# Add metadata back to splits
for i, split in enumerate(splits):
    split.metadata = {"source": f"chunk_{i}"}

# Create vector store
embeddings = OpenAIEmbeddings()
vectorstore = FAISS.from_documents(splits, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})

# Function to translate both company and investor names from English to Korean
def translate_entity_names(query):
    # Create a translation model
    translation_model = ChatOpenAI(model="gpt-4o-mini", temperature=0)
    
    # Updated prompt to handle both company and investor names
    translation_prompt = ChatPromptTemplate.from_messages([
        ("system", """You are a Korean-English translator specializing in Korean startup and investor names.
        Extract any potential company or investor names from the user's query and provide their Korean translations.
        
        If you detect company or investor names, return them in this format:
        "TRANSLATION_RESULT: {{English Name}} -> {{Korean Name}} (Type: company/investor)"
        
        You can return multiple translations if multiple entities are detected:
        "TRANSLATION_RESULT: 
        {{English Name 1}} -> {{Korean Name 1}} (Type: company)
        {{English Name 2}} -> {{Korean Name 2}} (Type: investor)"
        
        If no relevant names are detected, return:
        "TRANSLATION_RESULT: NO_ENTITIES_DETECTED"
        
        Do not provide any explanations or additional text, just the translation result."""),
        ("human", "{query}")
    ])
    
    translation_chain = translation_prompt | translation_model | StrOutputParser()
    result = translation_chain.invoke({"query": query})
    
    # Process the translation result
    translations = []
    modified_query = query
    
    if "TRANSLATION_RESULT:" in result:
        result_parts = result.strip().split("TRANSLATION_RESULT:")[1].strip()
        
        if result_parts != "NO_ENTITIES_DETECTED":
            # Handle multiple translations if present
            translation_lines = [line.strip() for line in result_parts.split('\n') if line.strip()]
            
            for line in translation_lines:
                if "->" in line:
                    # Extract parts
                    name_part, type_part = line.rsplit("(Type:", 1) if "(Type:" in line else (line, "(Type: unknown)")
                    english_name, korean_name = name_part.split("->")
                    entity_type = type_part.replace(")", "").strip()
                    
                    translation = {
                        "english_name": english_name.strip(),
                        "korean_name": korean_name.strip(),
                        "entity_type": entity_type.strip()
                    }
                    translations.append(translation)
                    
                    # Replace the English name with Korean name in the query
                    modified_query = modified_query.replace(english_name.strip(), korean_name.strip())
    
    if translations:
        return {
            "translated": True,
            "translations": translations,
            "modified_query": modified_query
        }
    else:
        return {
            "translated": False,
            "original_query": query
        }

# Create a system template that instructs the model how to handle the data
system_template = """You are an expert financial analyst specializing in Korean startup investments.
Answer the user's question based on the following retrieved information.
If you don't know the answer, say that you don't know.
Use the retrieved information to provide insights, summaries, and analysis.

Retrieved information:
{context}

Additional instructions:
- Translate any Korean text to English in your response
- If comparing multiple startups or sectors, create a structured comparison
- If analyzing trends, provide clear patterns and insights
- Be precise with numbers and statistics
- If information is missing, acknowledge it rather than making assumptions
- When referencing any company name, use both the English translated name and also the original Korean name in bracket
- Do the same for investor names - provide both English and Korean names where possible
"""

# Create a human template for the user's questions
human_template = "{question}"

# Create a chat prompt template
chat_prompt = ChatPromptTemplate.from_messages([
    ("system", system_template),
    ("human", human_template),
])

# Create a chat model
model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# Define the RAG chain
def retrieve_and_generate(query_info):
    # Extract the query and context from the input
    query = query_info["question"]
    context = query_info.get("context", "")
    
    # Run the chat model
    result = chat_prompt.invoke({
        "context": context, 
        "question": query
    })
    
    return model.invoke(result)

# Function to handle specialized queries
def handle_query(query):
    # First, attempt to translate any company or investor names in the query
    translation_result = translate_entity_names(query)
    
    # If translation was successful, use the modified query for retrieval
    if translation_result.get("translated", False):
        original_query = query
        query_for_retrieval = translation_result["modified_query"]
        
        # Add translation info to the query for the model
        translation_notes = []
        for t in translation_result["translations"]:
            translation_notes.append(f"'{t['english_name']}' to '{t['korean_name']}' (Type: {t['entity_type']})")
        
        translation_info = ", ".join(translation_notes)
        enhanced_query = f"{original_query}\n\nNote: I've translated: {translation_info} for better retrieval."
    else:
        query_for_retrieval = query
        enhanced_query = query
    
    # Check if query is about overall statistics
    if any(word in query.lower() for word in ["all", "overall", "summary", "statistics", "stats"]):
        # Generate statistical summary
        summary_stats = generate_statistics(df)
        
        # Add the statistics to the context
        context = f"Overall Dataset Statistics:\n{summary_stats}\n\n"
        
        # Use the retriever with the translated query
        retrieved_docs = retriever.invoke(query_for_retrieval)
        retrieved_context = "\n\n".join([doc.page_content for doc in retrieved_docs])
        
        combined_context = context + retrieved_context
        
        # Run the query with the statistics and retrieved context
        response = retrieve_and_generate({
            "context": combined_context,
            "question": enhanced_query
        })
        
        return StrOutputParser().invoke(response)
    
    # For all other queries, use the retriever with the translated query
    retrieved_docs = retriever.invoke(query_for_retrieval)
    retrieved_context = "\n\n".join([doc.page_content for doc in retrieved_docs])
    
    # Run the query with the retrieved context
    response = retrieve_and_generate({
        "context": retrieved_context,
        "question": enhanced_query
    })
    
    return StrOutputParser().invoke(response)

# Function to generate statistics from the DataFrame
def generate_statistics(df):
    stats = []
    
    # Count total number of startups
    stats.append(f"Total startups: {len(df)}")
    
    # Count unique sectors
    sectors = df["Sector"].dropna().unique()
    stats.append(f"Unique sectors: {len(sectors)}")
    
    # Count by funding rounds
    if "Round (KR)" in df.columns:
        round_counts = df["Round (KR)"].value_counts().to_dict()
        stats.append("Funding round distribution:")
        for round_type, count in round_counts.items():
            stats.append(f"  - {round_type}: {count}")
    
    # Calculate average investment if available
    if "Investment Amount" in df.columns:
        # This is simplified, need to properly parse the investment amounts
        stats.append("Investment statistics available in the dataset")
    
    return "\n".join(stats)

# Define the Gradio interface
demo = gr.Interface(
    fn=handle_query,
    inputs=gr.Textbox(label="Ask a question about Korean startup investments"),
    outputs=gr.Textbox(label="AI Answer"),
    title="Korean Startup Knowledge Base",
    description="Ask questions about Korean startups, funding trends, sectors, or investment rounds.",
    examples=[
        ["Which investors are most active in the healthcare sector?"],
        ["What was the funding trend for AI startups in 2025 so far?"],
        ["What do you know about Spoon Labs?"],
        ["Tell me about investments made by Kakao Investment."]
    ]
)

# Launch the Gradio interface
if __name__ == "__main__":
    demo.launch()