app.py

# import os
# import streamlit as st
# import torch
# from langchain.chains import LLMChain
# from langchain.prompts import ChatPromptTemplate
# from langchain_huggingface import HuggingFaceEndpoint

# def create_conversation_prompt(name1: str, name2: str, persona_style: str):
#     """
#     Create a prompt that instructs the model to produce exactly 15 messages
#     of conversation, alternating between name1 and name2, starting with name1.

#     We will be very explicit and not allow any formatting except the required lines.
#     """
#     prompt_template_str = f"""
#     You are simulating a conversation of exactly 15 messages between two people: {name1} and {name2}.
#     {name1} speaks first (message 1), then {name2} (message 2), then {name1} (message 3), and so forth,
#     alternating until all 15 messages are complete. The 15th message is by {name1}.

#     Requirements:
#     - Output exactly 15 lines, no more, no less.
#     - Each line must be a single message in the format:
#       {name1}: <message> or {name2}: <message>
#     - Do not add any headings, numbers, sample outputs, or explanations.
#     - Do not mention code, programming, or instructions.
#     - Each message should be 1-2 short sentences, friendly, natural, reflecting the style: {persona_style}.
#     - Use everyday language, can ask questions, show opinions.
#     - Use emojis sparingly if it fits the style (no more than 1-2 total).
#     - No repeated lines, each message should logically follow from the previous one.
#     - Do not produce anything after the 15th message. No extra lines or text.

#     Produce all 15 messages now:
#     """
#     return ChatPromptTemplate.from_template(prompt_template_str)

# def create_summary_prompt(name1: str, name2: str, conversation: str):
#     """Prompt for generating a title and summary."""
#     summary_prompt_str = f"""
#     Below is a completed 15-message conversation between {name1} and {name2}:

#     {conversation}

#     Please provide:
#     Title: <A short descriptive title of the conversation>
#     Summary: <A few short sentences highlighting the main points, tone, and conclusion>

#     Do not continue the conversation, do not repeat it, and do not add extra formatting beyond the two lines:
#     - One line starting with "Title:"
#     - One line starting with "Summary:"
#     """
#     return ChatPromptTemplate.from_template(summary_prompt_str)

# def main():
#     st.title("LLM Conversation Simulation")

#     model_names = [
#         "meta-llama/Llama-3.3-70B-Instruct",
#         "meta-llama/Llama-3.1-405B-Instruct",
#         "Qwen/Qwen2.5-72B-Instruct",
#         "deepseek-ai/DeepSeek-V3",
#         "deepseek-ai/DeepSeek-V2.5"
        
#     ]
#     selected_model = st.selectbox("Select a model:", model_names)

#     name1 = st.text_input("Enter the first user's name:", value="Alice")
#     name2 = st.text_input("Enter the second user's name:", value="Bob")
#     persona_style = st.text_area("Enter the persona style characteristics:", 
#                                  value="friendly, curious, and a bit sarcastic")

#     if st.button("Start Conversation Simulation"):
#         st.write("**Loading model...**")
#         print("Loading model...")

#         with st.spinner("Starting simulation..."):
#             endpoint_url = f"https://api-inference.huggingface.co/models/{selected_model}"

#             try:
#                 llm = HuggingFaceEndpoint(
#                     endpoint_url=endpoint_url,
#                     huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN"),
#                     task="text-generation",
#                     temperature=0.7,
#                     max_new_tokens=512
#                 )
#                 st.write("**Model loaded successfully!**")
#                 print("Model loaded successfully!")
#             except Exception as e:
#                 st.error(f"Error initializing HuggingFaceEndpoint: {e}")
#                 print(f"Error initializing HuggingFaceEndpoint: {e}")
#                 return

#             conversation_prompt = create_conversation_prompt(name1, name2, persona_style)
#             conversation_chain = LLMChain(llm=llm, prompt=conversation_prompt)

#             st.write("**Generating the full 15-message conversation...**")
#             print("Generating the full 15-message conversation...")

#             try:
#                 # Generate all 15 messages in one go
#                 conversation = conversation_chain.run(chat_history="", input="").strip()

#                 st.subheader("Final Conversation:")
#                 st.text(conversation)
#                 print("Conversation Generation Complete.\n")
#                 print("Full Conversation:\n", conversation)

#                 # Summarize the conversation
#                 summary_prompt = create_summary_prompt(name1, name2, conversation)
#                 summary_chain = LLMChain(llm=llm, prompt=summary_prompt)

#                 st.subheader("Summary and Title:")
#                 st.write("**Summarizing the conversation...**")
#                 print("Summarizing the conversation...")

#                 summary = summary_chain.run(chat_history="", input="")
#                 st.write(summary)
#                 print("Summary:\n", summary)

#             except Exception as e:
#                 st.error(f"Error generating conversation: {e}")
#                 print(f"Error generating conversation: {e}")

# if __name__ == "__main__":
#     main()


import os
import streamlit as st
import google.cloud.aiplatform as aiplatform
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate
from langchain.llms.base import LLM
from pydantic import BaseModel
from typing import Optional, List, Mapping, Any

###############################################################################
# 1. Create a Custom LLM class for LangChain to call your Vertex AI endpoint.
###############################################################################
class VertexAICustomModel(LLM, BaseModel):
    project_id: str
    location: str
    endpoint_id: str
    temperature: float = 0.7
    max_new_tokens: int = 512
    
    @property
    def _llm_type(self) -> str:
        return "vertex_ai_custom"

    def _call(self, prompt: str, stop: Optional[List[str]] = None) -> str:
        # Initialize Vertex AI with your project/region
        aiplatform.init(project=self.project_id, location=self.location)
        endpoint = aiplatform.Endpoint(
            endpoint_name=f"projects/{self.project_id}/locations/{self.location}/endpoints/{self.endpoint_id}"
        )
        
        # Construct the instance for prediction. 
        # NOTE: Adjust 'prompt', 'temperature', etc. if your model expects different parameters.
        instance = {
            "prompt": prompt,
            "temperature": self.temperature,
            "max_new_tokens": self.max_new_tokens
        }

        # Call the endpoint
        response = endpoint.predict(instances=[instance])
        
        # Extract the text from the response.
        # This will vary depending on how your model returns predictions.
        # A common approach is response.predictions[0]["generated_text"], 
        # but confirm your model's actual JSON structure.
        predictions = response.predictions
        if not predictions or "generated_text" not in predictions[0]:
            raise ValueError(
                f"Unexpected response structure from Vertex AI endpoint: {response}"
            )
        
        text = predictions[0]["generated_text"]

        # Optionally apply 'stop' tokens
        if stop:
            for s in stop:
                if s in text:
                    text = text.split(s)[0]
        return text

    @property
    def _identifying_params(self) -> Mapping[str, Any]:
        """Return any identifying parameters of this LLM."""
        return {
            "endpoint_id": self.endpoint_id,
            "project_id": self.project_id,
            "location": self.location,
            "temperature": self.temperature,
            "max_new_tokens": self.max_new_tokens,
        }

###############################################################################
# 2. Create your conversation and summary prompt templates (unchanged).
###############################################################################
def create_conversation_prompt(name1: str, name2: str, persona_style: str):
    """
    Create a prompt that instructs the model to produce exactly 15 messages
    of conversation, alternating between name1 and name2, starting with name1.
    """
    prompt_template_str = f"""
    You are simulating a conversation of exactly 15 messages between two people: {name1} and {name2}.
    {name1} speaks first (message 1), then {name2} (message 2), then {name1} (message 3), and so forth,
    alternating until all 15 messages are complete. The 15th message is by {name1}.

    Requirements:
    - Output exactly 15 lines, no more, no less.
    - Each line must be a single message in the format:
      {name1}: <message> or {name2}: <message>
    - Do not add any headings, numbers, sample outputs, or explanations.
    - Do not mention code, programming, or instructions.
    - Each message should be 1-2 short sentences, friendly, natural, reflecting the style: {persona_style}.
    - Use everyday language, can ask questions, show opinions.
    - Use emojis sparingly if it fits the style (no more than 1-2 total).
    - No repeated lines, each message should logically follow from the previous one.
    - Do not produce anything after the 15th message. No extra lines or text.

    Produce all 15 messages now:
    """
    return ChatPromptTemplate.from_template(prompt_template_str)

def create_summary_prompt(name1: str, name2: str, conversation: str):
    """Prompt for generating a title and summary."""
    summary_prompt_str = f"""
    Below is a completed 15-message conversation between {name1} and {name2}:
    {conversation}
    Please provide:
    Title: <A short descriptive title of the conversation>
    Summary: <A few short sentences highlighting the main points, tone, and conclusion>
    Do not continue the conversation, do not repeat it, and do not add extra formatting beyond the two lines:
    - One line starting with "Title:"
    - One line starting with "Summary:"
    """
    return ChatPromptTemplate.from_template(summary_prompt_str)

###############################################################################
# 3. Main Streamlit app with Vertex AI usage.
###############################################################################
def main():
    st.title("LLM Conversation Simulation (GCP Vertex AI)")

    # We can remove model selection if we are always using your deployed model:
    # st.selectbox(... ) # => Removed

    # Hardcode or load your Vertex AI endpoint details here
    project_id = "282802344966"
    location = "us-west1"
    endpoint_id = "1106913540054188032"

    # Input fields for conversation
    name1 = st.text_input("Enter the first user's name:", value="Alice")
    name2 = st.text_input("Enter the second user's name:", value="Bob")
    persona_style = st.text_area("Enter the persona style characteristics:",
                                 value="friendly, curious, and a bit sarcastic")

    if st.button("Start Conversation Simulation"):
        st.write("**Initializing Vertex AI endpoint...**")
        st.spinner("Starting simulation...")

        # Create your custom LLM that calls Vertex AI
        llm = VertexAICustomModel(
            project_id=project_id,
            location=location,
            endpoint_id=endpoint_id,
            temperature=0.7,
            max_new_tokens=512
        )

        st.write("**Vertex AI endpoint loaded successfully!**")

        # Build the conversation chain
        conversation_prompt = create_conversation_prompt(name1, name2, persona_style)
        conversation_chain = LLMChain(llm=llm, prompt=conversation_prompt)

        st.write("**Generating the full 15-message conversation...**")

        try:
            # Generate all 15 messages in one go
            conversation = conversation_chain.run(chat_history="", input="").strip()

            st.subheader("Final Conversation:")
            st.text(conversation)

            # Summarize the conversation
            summary_prompt = create_summary_prompt(name1, name2, conversation)
            summary_chain = LLMChain(llm=llm, prompt=summary_prompt)

            st.subheader("Summary and Title:")
            st.write("**Summarizing the conversation...**")

            summary = summary_chain.run(chat_history="", input="")
            st.write(summary)

        except Exception as e:
            st.error(f"Error generating conversation: {e}")
            print(f"Error generating conversation: {e}")

if __name__ == "__main__":
    main()