demo_app.py

import streamlit as st
import json
import time
from client import VcGeminiV0Env
from models import VcGeminiV0Action
import pandas as pd
import matplotlib.pyplot as plt

# Premium Page Config
st.set_page_config(
    page_title="VC Gemini V0 - Live Simulation",
    page_icon="💸",
    layout="wide",
)

# Dark theme CSS
st.markdown("""
<style>
    .reportview-container {
        background: #0e1117;
    }
    .main {
        background: #0e1117;
        color: #ffffff;
    }
    .stButton>button {
        width: 100%;
        border-radius: 5px;
        height: 3em;
        background-color: #4CAF50;
        color: white;
    }
    .stMetric {
        background-color: #1e2130 !important;
        padding: 15px !important;
        border-radius: 10px !important;
        border: 1px solid #30363d !important;
    }
    div[data-testid="stMetricValue"] {
        color: #4CAF50 !important;
    }
    div[data-testid="stMetricLabel"] {
        color: #ffffff !important;
    }
</style>
""", unsafe_allow_html=True)

st.title("💸 VC Gemini V0: AI Investment Agent")
st.markdown("*A state-of-the-art RL agent evaluating startups in real-time.*")

# Sidebar - Configuration
with st.sidebar:
    st.header("Settings")
    env_url = st.text_input("Environment URL", value="https://shrads78-vc-gemini-v0.hf.space")
    model_type = st.selectbox("LLM Strategy", ["Llama-3.1-70B (Baseline)", "Qwen-2.5-3B (GRPO Trained)"])
    if st.button("Reset Simulation"):
        st.session_state.obs = None
        st.session_state.history = []
        st.rerun()

# Session State Initialization
if "obs" not in st.session_state:
    st.session_state.obs = None
    st.session_state.history = []
    st.session_state.env_client = None

def get_client():
    if st.session_state.env_client is None:
        st.session_state.env_client = VcGeminiV0Env(base_url=env_url)
    return st.session_state.env_client

# Main UI Columns
col1, col2 = st.columns([2, 1])

with col1:
    st.header("Narrative Log")
    log_container = st.container()
    
    if st.session_state.obs is None:
        if st.button("Start New Fund (Reset)"):
            client = get_client()
            with st.spinner("Connecting to environment..."):
                obs = client.reset()
                st.session_state.obs = obs
                obs_text = obs.observation.observation_text if hasattr(obs, 'observation') else obs.observation_text
                st.session_state.history.append({"role": "system", "content": obs_text})
                st.rerun()
    
    # Display History
    for entry in st.session_state.history:
        if entry["role"] == "system":
            st.info(entry["content"])
        else:
            st.success(f"**Action:** {entry['content']}")

with col2:
    st.header("Portfolio Stats")
    if st.session_state.obs:
        data = st.session_state.obs.observation.data
        budget = data.get("budget", 100000000.0)
        quarter = data.get("quarter", 1)
        turns_left = data.get("turns_left", 5)
        
        st.metric("Quarter", f"Q{quarter}")
        st.metric("Fund Budget", f"${budget:,.0f}")
        st.metric("Turns Left", turns_left)
        
        # Portfolio visualization
        portfolio = data.get("portfolio", [])
        if portfolio:
            df = pd.DataFrame(portfolio)
            st.subheader("Holdings")
            st.table(df[["startup_name", "invested_amount", "paper_multiplier"]])
            
            # Simple Chart
            fig, ax = plt.subplots()
            ax.pie(df['invested_amount'], labels=df['startup_name'], autopct='%1.1f%%', startangle=90)
            ax.set_title("Capital Allocation")
            st.pyplot(fig)

# Actions
if st.session_state.obs and not st.session_state.obs.done:
    st.divider()
    st.header("Take Action")
    act_col1, act_col2, act_col3 = st.columns(3)
    
    with act_col1:
        st.info("💡 **Hint:** Check 'Quarter 1' or 'Market Updates' in Wait/Read to discover startup names (e.g. 'Aura Robotics', 'QuantumCore').")
        startup_to_invest = st.text_input("Startup Name to Invest In", placeholder="e.g. Aura Robotics")
        if st.button("Invest $20M"):
            client = get_client()
            action = VcGeminiV0Action(action_type="invest", parameters={"startup_name": startup_to_invest})
            res = client.step(action)
            st.session_state.obs = res
            st.session_state.history.append({"role": "user", "content": f"Invest in {startup_to_invest}"})
            obs_text = res.observation.observation_text if hasattr(res, 'observation') else getattr(res, 'observation_text', 'No observation text')
            st.session_state.history.append({"role": "system", "content": obs_text})
            st.rerun()
            
    with act_col2:
        st.info("💡 **Hint:** Read market overview by querying 'Quarter 1' or read a startup's deck, e.g., 'Aura_Robotics/pitch_deck.txt'.")
        path_to_read = st.text_input("Path to Diligence File", placeholder="Aura_Robotics/pitch_deck.txt")
        if st.button("Read File"):
            client = get_client()
            action = VcGeminiV0Action(action_type="read_file", parameters={"path": path_to_read})
            res = client.step(action)
            st.session_state.obs = res
            st.session_state.history.append({"role": "user", "content": f"Read {path_to_read}"})
            obs_text = res.observation.observation_text if hasattr(res, 'observation') else getattr(res, 'observation_text', 'No observation text')
            st.session_state.history.append({"role": "system", "content": obs_text})
            st.rerun()

    with act_col3:
        st.info("💡 **Hint:** Wait a turn to let time advance when you have no other diligence actions.")
        if st.button("Wait a Turn"):
            client = get_client()
            action = VcGeminiV0Action(action_type="wait")
            res = client.step(action)
            st.session_state.obs = res
            st.session_state.history.append({"role": "user", "content": "Wait"})
            obs_text = res.observation.observation_text if hasattr(res, 'observation') else getattr(res, 'observation_text', 'No observation text')
            st.session_state.history.append({"role": "system", "content": obs_text})
            st.rerun()

if st.session_state.obs and st.session_state.obs.done:
    st.balloons()
    st.header("Fund Performance Results")
    final_tvpi = st.session_state.obs.observation.data.get("tvpi", 0.0)
    st.metric("FINAL TVPI", f"{final_tvpi:.2f}x")
    if final_tvpi >= 1.20:
        st.success("SUCCESS! You beat the Hurdle Rate.")
    else:
        st.error("FAILURE. Below Hurdle Rate.")