app.py

import os
import streamlit as st
from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenerativeAIEmbeddings, ChatGoogleGenerativeAI
from langchain_classic.chains import RetrievalQA

# 1. Load the API Key safely from the local .env file
with open('.env') as f:
    for line in f:
        if '=' in line:
            key, value = line.strip().split('=', 1)
            os.environ[key] = value.strip('"\'')

st.title("🛠️ API Breaking-Change Healer")
st.write("Upload your broken legacy code, and I'll rewrite it using the latest docs!")

# 2. Build Vector Database (Cached so it doesn't reload every time you click a button)
@st.cache_resource
def build_vector_store():
    # Load our mock migration guide
    loader = TextLoader("migration_guide.txt")
    docs = loader.load()
    
    # Split the document into readable chunks
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
    splits = text_splitter.split_documents(docs)
    
    # Create the Vector Database using the active Gemini embedding model
    embeddings = GoogleGenerativeAIEmbeddings(model="models/gemini-embedding-001")
    vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
    
    return vectorstore

# Initialize the database
db = build_vector_store()
st.success("Migration guide loaded into memory! Ready to heal some code.")

# 3. The RAG UI and Logic
st.subheader("Paste your broken legacy code below:")
broken_code = st.text_area("Old Code:", height=150)

if st.button("Heal My Code! ✨"):
    if broken_code:
        with st.spinner("Consulting the migration guide..."):
            
            # Initialize the LLM using the ultra-stable 'gemini-pro' model
            llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash", temperature=0)
            
            # Create the RAG chain linking the Database to the LLM
            qa_chain = RetrievalQA.from_chain_type(
                llm=llm,
                chain_type="stuff",
                retriever=db.as_retriever(search_kwargs={"k": 2})
            )
            
            # Formulate the prompt with the exact instructions
            prompt = f"""
            You are an expert developer helping to migrate old JavaScript code to a new API version.
            Use ONLY the provided migration guide context to fix the following broken code.
            Output the fixed code in a clean markdown block, followed by a brief bulleted list of what you changed.
            
            Broken Code:
            {broken_code}
            """
            
            # Execute the RAG pipeline
            result = qa_chain.invoke(prompt)
            
            # Display results
            st.markdown("### 🩹 Fixed Code & Explanation")
            st.write(result["result"])
            st.balloons()
    else:
        st.warning("Please paste some code first!")