Update app.py

app.py

@@ -1,16 +1,208 @@
+# 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 torch
+import google.cloud.aiplatform as aiplatform
 from langchain.chains import LLMChain
 from langchain.prompts import ChatPromptTemplate
-from langchain_huggingface import HuggingFaceEndpoint
+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.
-
-    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}.
@@ -37,90 +229,77 @@ 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")
-
-    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)
+    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:", 
+    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}")
+        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()
-