Update neo_sages5.py

neo_sages5.py

@@ -1,627 +1,446 @@
 import streamlit as st
-from io import BytesIO
-import ibm_watsonx_ai
-import secretsload
-import genparam
-import requests
-import time
-import re
-import json
-import cos_creds
-from ibm_watsonx_ai.foundation_models import ModelInference
-from ibm_watsonx_ai import Credentials, APIClient
-from ibm_watsonx_ai.metanames import GenTextParamsMetaNames as GenParams
-from ibm_watsonx_ai.metanames import GenTextReturnOptMetaNames as RetParams
-from knowledge_bases import KNOWLEDGE_BASE_OPTIONS, SYSTEM_PROMPTS, VECTOR_INDEXES
-from ibm_watsonx_ai.foundation_models import Embeddings
-from ibm_watsonx_ai.foundation_models.utils.enums import EmbeddingTypes
-from pymilvus import MilvusClient
-
-from secretsload import load_stsecrets
-
-credentials = load_stsecrets()
-
-st.set_page_config(
-    page_title="The Solutioning Sages",
-    page_icon="🪄",
-    initial_sidebar_state="collapsed",
-    layout="wide"
-)
-
-# Password protection
-def check_password():
-    def password_entered():
-        if st.session_state["password"] == st.secrets["app_password"]:
-            st.session_state["password_correct"] = True
-            del st.session_state["password"]
-        else:
-            st.session_state["password_correct"] = False
-
-    if "password_correct" not in st.session_state:
-        st.markdown("\n\n")
-        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
-        st.divider()
-        st.info("Designed and developed by Milan Mrdenovic © IBM Norway 2024")
-        return False
-    elif not st.session_state["password_correct"]:
-        st.markdown("\n\n")
-        st.text_input("Enter the password", type="password", on_change=password_entered, key="password")
-        st.divider()
-        st.error("😕 Incorrect password")
-        st.info("Designed and developed by Milan Mrdenovic © IBM Norway 2024")
-        return False
-    else:
-        return True
-
-def initialize_session_state():
-    if 'chat_history_1' not in st.session_state:
-        st.session_state.chat_history_1 = []
-    if 'chat_history_2' not in st.session_state:
-        st.session_state.chat_history_2 = []
-    if 'chat_history_3' not in st.session_state:
-        st.session_state.chat_history_3 = []
-    if 'first_question' not in st.session_state:
-        st.session_state.first_question = False 
-    if "counter" not in st.session_state:
-        st.session_state["counter"] = 0
-    if 'token_statistics' not in st.session_state:
-        st.session_state.token_statistics = []
-    if 'selected_kb' not in st.session_state:
-        st.session_state.selected_kb = KNOWLEDGE_BASE_OPTIONS[0]
-    if 'current_system_prompts' not in st.session_state:
-        st.session_state.current_system_prompts = SYSTEM_PROMPTS[st.session_state.selected_kb]
-
-# three_column_style = """
-#     <style>
-#     .stColumn {
-#         padding: 0.5rem;
-#         border-right: 1px solid #dedede;
-#     }
-#     .stColumn:last-child {
-#         border-right: none;
-#     }
-#     .chat-container {
-#         height: calc(100vh - 200px);
-#         overflow-y: auto;
-#     }
-#     </style>
-# """
-
-three_column_style = """
-    <style>
-    .stColumn {
-        padding: 0.5rem;
-        border-right: 1px solid #dedede;
-    }
-    .stColumn:last-child {
-        border-right: none;
-    }
-    .chat-container {
-        height: calc(100vh - 200px);
-        overflow-y: auto;
-        display: flex;
-        flex-direction: column;
-    }
-    .chat-messages {
-        display: flex;
-        flex-direction: column;
-        gap: 1rem;
-    }
-    </style>
-""" # Alt Style
-
-#-----
-def get_active_model():
-    return genparam.SELECTED_MODEL_1 if genparam.ACTIVE_MODEL == 0 else genparam.SELECTED_MODEL_2
-
-def get_active_prompt_template():
-    return genparam.PROMPT_TEMPLATE_1 if genparam.ACTIVE_MODEL == 0 else genparam.PROMPT_TEMPLATE_2
-
-def get_active_vector_index():
-    selected_kb = st.session_state.selected_kb
-    if genparam.ACTIVE_INDEX == 0:
-        return VECTOR_INDEXES[selected_kb]["index_1"]
-    else:
-        return VECTOR_INDEXES[selected_kb]["index_2"]
-#-----
+from langchain import PromptTemplate
+from typing import TypedDict, List, Dict, Optional
+from langchain.graphs import StateGraph
+from dataclasses import dataclass, field
+import random
+
+# Data Structures
+@dataclass
+class StoryState:
+    current_step: int = 0
+    max_steps: int = 5
+    story_log: List[str] = field(default_factory=list)
+    user_inputs: List[str] = field(default_factory=list)
+    character1_responses: List[str] = field(default_factory=list)
+    character2_responses: List[str] = field(default_factory=list)
+    story_outcome: Optional[str] = None
+
+class MicroStory(TypedDict):
+    title: str
+    initial_setup: str
+    character1_name: str
+    character2_name: str
+    steps: List[str]
+    success_conditions: List[str]
+    failure_conditions: List[str]
+
+from langchain.graphs import StateGraph
+from typing import Dict, List, Any
+from dataclasses import dataclass
+from enum import Enum
+
+class StoryNodeType(Enum):
+    SETUP = "setup"
+    USER_INPUT = "user_input"
+    CHARACTER1_RESPONSE = "character1_response"
+    CHARACTER2_RESPONSE = "character2_response"
+    EVALUATION = "evaluation"
+
+@dataclass
+class StoryGraphState:
+    current_node: StoryNodeType
+    story_data: Dict[str, Any]
+    accumulated_context: List[Dict[str, str]]
+    step_count: int = 0
+
+def create_story_graph() -> StateGraph:
+    """Creates the state graph for story progression"""
+    graph = StateGraph()
+
+    # Define state transitions
+    def setup_to_user_input(state: StoryGraphState) -> StoryGraphState:
+        state.current_node = StoryNodeType.USER_INPUT
+        return state
+
+    def user_input_to_char1(state: StoryGraphState, user_input: str) -> StoryGraphState:
+        state.current_node = StoryNodeType.CHARACTER1_RESPONSE
+        state.accumulated_context.append({"role": "user", "content": user_input})
+        return state
+
+    def char1_to_char2(state: StoryGraphState, char1_response: str) -> StoryGraphState:
+        state.current_node = StoryNodeType.CHARACTER2_RESPONSE
+        state.accumulated_context.append({"role": "character1", "content": char1_response})
+        return state
+
+    def char2_to_evaluation(state: StoryGraphState, char2_response: str) -> StoryGraphState:
+        state.current_node = StoryNodeType.EVALUATION
+        state.accumulated_context.append({"role": "character2", "content": char2_response})
+        state.step_count += 1
+        return state
+
+    def evaluation_to_next(state: StoryGraphState) -> StoryGraphState:
+        if state.step_count >= 5:
+            # Story is complete, stay in evaluation
+            return state
+        # Move to next user input
+        state.current_node = StoryNodeType.USER_INPUT
+        return state
+
+    # Add nodes and edges
+    graph.add_node("setup", setup_to_user_input)
+    graph.add_node("user_input", user_input_to_char1)
+    graph.add_node("character1_response", char1_to_char2)
+    graph.add_node("character2_response", char2_to_evaluation)
+    graph.add_node("evaluation", evaluation_to_next)
+
+    # Connect nodes
+    graph.add_edge("setup", "user_input")
+    graph.add_edge("user_input", "character1_response")
+    graph.add_edge("character1_response", "character2_response")
+    graph.add_edge("character2_response", "evaluation")
+    graph.add_edge("evaluation", "user_input")
+
+    return graph
+
+class StoryRunner:
+    def __init__(self, story_data: Dict[str, Any]):
+        self.graph = create_story_graph()
+        self.state = StoryGraphState(
+            current_node=StoryNodeType.SETUP,
+            story_data=story_data,
+            accumulated_context=[]
+        )
 
-def setup_client(project_id=None):
-    credentials = Credentials(
-        url=st.secrets["url"],
-        api_key=st.secrets["api_key"]
-    )
-    # Use the passed project_id if provided, otherwise fallback to secrets
-    project_id = project_id or st.secrets["project_id"]
-    client = APIClient(credentials, project_id=project_id)
-    return credentials, client
-
-wml_credentials, client = setup_client(st.secrets["project_id"])
-
-import ibm_boto3
-from ibm_botocore.client import Config
-###==========================================1
-def setup_retrieval_cos_client(config_dict=None):
-    # Default credentials (same as main notebook for now)
-    default_config = {
-        "api_key": cos_creds.COS_API_KEY,
-        "instance_id": cos_creds.COS_INSTANCE_ID,
-        "endpoint": cos_creds.COS_ENDPOINT,
-        "bucket": cos_creds.BUCKET_NAME
-    }
-    
-    # Use provided config or default
-    config = config_dict if config_dict is not None else default_config
-    
-    # Initialize the retrieval COS client
-    retrieval_cos_client = ibm_boto3.client(
-        "s3",
-        ibm_api_key_id=config["api_key"],
-        ibm_service_instance_id=config["instance_id"],
-        config=Config(signature_version="oauth"),
-        endpoint_url=config["endpoint"]
-    )
-    
-    # Verify the connection by trying to list objects
-    try:
-        retrieval_cos_client.list_objects(Bucket=config["bucket"], MaxKeys=1)
-        print("Retrieval COS client successfully initialized and connected")
-    except Exception as e:
-        print(f"Error verifying retrieval COS client connection: {str(e)}")
-        raise
-    
-    return retrieval_cos_client
-
-retrieval_cos_client = setup_retrieval_cos_client()
-
-def load_callable_index_config(config_path):
-    try:
-        # Download config file content using retrieval client
-        response = retrieval_cos_client.get_object(Bucket=cos_creds.BUCKET_NAME, Key=config_path)
-        config_content = response['Body'].read().decode('utf-8')
-        return json.loads(config_content)
-    except Exception as e:
-        raise Exception(f"Error loading callable index config: {str(e)}")
-###==========================================1-2
-def setup_vector_index(client, wml_credentials, config_path):
-    # Load the configuration using load_callable_index_config
-    config = load_callable_index_config(config_path)
-    
-    # Initialize embeddings
-    emb = Embeddings(
-        model_id=config["embedding"]["model_id"],
-        credentials=wml_credentials,
-        project_id=PROJECT_ID,
-        params={
-            "truncate_input_tokens": config["embedding"]["max_tokens"]
+    def process_user_input(self, user_input: str) -> Dict[str, Any]:
+        """Process user input and advance the story state"""
+        if self.state.current_node != StoryNodeType.USER_INPUT:
+            raise ValueError("Not ready for user input")
+
+        # Advance through the graph
+        self.state = self.graph.transition("user_input", self.state, user_input)
+        self.state = self.graph.transition("character1_response", self.state, None)
+        self.state = self.graph.transition("character2_response", self.state, None)
+        self.state = self.graph.transition("evaluation", self.state, None)
+
+        # Return current state info
+        return {
+            "step_count": self.state.step_count,
+            "is_complete": self.state.step_count >= 5,
+            "current_context": self.state.accumulated_context[-3:] if self.state.accumulated_context else [],
+            "current_node": self.state.current_node.value
         }
-    )
-    
-    # Get connection details
-    connection_details = client.connections.get_details(config["connection"]["connection_id"])
-    connection_properties = connection_details["entity"]["properties"]
-    
-    # Initialize Milvus client
-    milvus_client = MilvusClient(
-        uri=f'https://{connection_properties.get("host")}:{connection_properties.get("port")}',
-        user=connection_properties.get("username"),
-        password=connection_properties.get("password"),
-        db_name=config["connection"]["database"]
-    )
-    
-    # Prepare vector index properties
-    vector_index_properties = {
-        "store": {
-            "index": config["collection"]["name"],
-            "database": config["connection"]["database"]
+
+    def get_full_context(self) -> List[Dict[str, str]]:
+        """Get the full conversation context"""
+        return self.state.accumulated_context
+
+    def is_complete(self) -> bool:
+        """Check if the story is complete"""
+        return self.state.step_count >= 5
+
+# Sample Stories Database
+# Story Categories and Templates
+STORY_CATEGORIES = {
+    "Mystery": [
+        {
+            "title": "The Library Mystery",
+            "initial_setup": "In the ancient library of St. Bartholomew's, a rare manuscript has gone missing.",
+            "character1_name": "Detective Nash",
+            "character2_name": "Librarian Wells",
+            "steps": [
+                "You notice strange symbols carved into the reading desk",
+                "A student mentions seeing someone in medieval clothing",
+                "The manuscript tracking system shows impossible timestamps",
+                "Temperature drops significantly in the rare books section",
+                "You find a hidden door behind the card catalog"
+            ],
+            "success_conditions": [
+                "mentioned checking the security cameras",
+                "investigated the symbols",
+                "questioned the student further",
+                "connected medieval sighting with timestamps"
+            ],
+            "failure_conditions": [
+                "accused the librarian",
+                "ignored the symbols",
+                "left the library",
+                "called the police immediately"
+            ]
         },
-        "settings": {
-            "embedding_model_id": config["embedding"]["model_id"],
-            "schema_fields": config["collection"]["schema"]["fields"],
-            "top_k": config["index_settings"]["top_k"]
+        {
+            "title": "The Digital Deception",
+            "initial_setup": "A tech startup's revolutionary AI algorithm has been stolen right before a major demo.",
+            "character1_name": "Cyber Detective Chen",
+            "character2_name": "System Admin Rodriguez",
+            "steps": [
+                "The server logs show multiple failed login attempts",
+                "An employee reports receiving a suspicious email",
+                "The backup system was manually disabled",
+                "Strange network traffic appears during off-hours",
+                "A hidden backdoor program is discovered"
+            ],
+            "success_conditions": [
+                "checked email headers",
+                "analyzed network logs",
+                "investigated backup system",
+                "traced the backdoor"
+            ],
+            "failure_conditions": [
+                "restored from backup immediately",
+                "ignored the suspicious email",
+                "reset all passwords without investigation",
+                "blamed the system admin"
+            ]
         }
-    }
-    
-    # Return the entire config instead of just the schema fields
-    return milvus_client, emb, vector_index_properties, config
+    ],
+    "Adventure": [
+        {
+            "title": "The Lost Temple",
+            "initial_setup": "Deep in the Amazon rainforest, you've discovered the entrance to an ancient temple.",
+            "character1_name": "Dr. Rivera",
+            "character2_name": "Guide Santos",
+            "steps": [
+                "Ancient markings warn of a curse",
+                "You find a mechanism with multiple levers",
+                "A strange humming sound emanates from deeper within",
+                "The floor tiles show a peculiar pattern",
+                "A beam of light reveals a hidden chamber"
+            ],
+            "success_conditions": [
+                "documented the markings",
+                "observed the pattern",
+                "tested the mechanism carefully",
+                "followed the light beam"
+            ],
+            "failure_conditions": [
+                "ignored the warnings",
+                "pulled levers randomly",
+                "split up the group",
+                "took artifacts without examination"
+            ]
+        }
+    ],
+    "Sci-Fi": [
+        {
+            "title": "The Quantum Anomaly",
+            "initial_setup": "At a cutting-edge research facility, a quantum experiment has created an unexplained phenomenon.",
+            "character1_name": "Dr. Zhang",
+            "character2_name": "Engineer Parker",
+            "steps": [
+                "Quantum readings are off the charts",
+                "Equipment starts behaving erratically",
+                "A shimmer appears in the air",
+                "Time seems to flow differently near the anomaly",
+                "Multiple reality signatures detected"
+            ],
+            "success_conditions": [
+                "monitored quantum fluctuations",
+                "calibrated equipment",
+                "documented time discrepancies",
+                "maintained safe distance"
+            ],
+            "failure_conditions": [
+                "shut down power immediately",
+                "entered the anomaly",
+                "ignored safety protocols",
+                "attempted to contain without data"
+            ]
+        }
+    ]
+}
+
+# Flatten categories for easy access by title
+STORY_LOOKUP = {
+    story["title"]: story 
+    for category in STORY_CATEGORIES.values() 
+    for story in category
+}
+
+# Character Response Templates
+CHARACTER1_TEMPLATE = """
+Context: You are {character1_name} in this story.
+Story Progress: {story_log}
+User's Latest Action: {user_input}
+
+Respond to the user's action in character, considering:
+1. Your role and personality
+2. The current story situation
+3. The potential consequences of their action
+
+Response:
+"""
+
+CHARACTER2_TEMPLATE = """
+Context: You are {character2_name} in this story.
+Story Progress: {story_log}
+User's Latest Action: {user_input}
+Other Character's Response: {character1_response}
+
+Respond to both the user and {character1_name}, considering:
+1. Your role and personality
+2. The current story developments
+3. Your relationship with {character1_name}
+4. The potential impact on the story's outcome
+
+Response:
+"""
 
-###==========================================1-2
-    def proximity_search(question, milvus_client, emb, vector_index_properties, config):
+def initialize_session_state():
+    if 'story_state' not in st.session_state:
+        st.session_state.story_state = StoryState()
+    if 'selected_category' not in st.session_state:
+        st.session_state.selected_category = list(STORY_CATEGORIES.keys())[0]
+    if 'current_story' not in st.session_state:
+        # Select random story from current category
+        st.session_state.current_story = random.choice(STORY_CATEGORIES[st.session_state.selected_category])
+
+def evaluate_outcome(state: StoryState, story: MicroStory) -> str:
+    user_actions = " ".join(state.user_inputs).lower()
     
-    query_vectors = emb.embed_query(question)
-    schema_fields = config["collection"]["schema"]["fields"]
-    milvus_response = milvus_client.search(
-        collection_name=vector_index_properties["store"]["index"],
-        data=[query_vectors],
-        limit=vector_index_properties["settings"]["top_k"],
-        metric_type="L2",
-        output_fields=[
-            schema_fields["text_field"],
-            schema_fields["document_name"],
-            schema_fields["page_number"]
-        ]
-    )
-
-    documents = []
+    # Count matches for success and failure conditions
+    success_matches = sum(1 for cond in story["success_conditions"] if cond.lower() in user_actions)
+    failure_matches = sum(1 for cond in story["failure_conditions"] if cond.lower() in user_actions)
     
-    for hit in milvus_response[0]:
-        text = hit["entity"].get(schema_fields["text_field"], "")
-        doc_name = hit["entity"].get(schema_fields["document_name"], "Unknown Document")
-        page_num = hit["entity"].get(schema_fields["page_number"], "N/A")
-        
-        formatted_result = f"Document: {doc_name}\nContent: {text}\nPage: {page_num}\n"
-        documents.append(formatted_result)
+    # Calculate success ratio
+    success_ratio = success_matches / len(story["success_conditions"])
     
-    # Format final output
-    joined = "\n".join(documents)
-    retrieved = f"Number of Retrieved Documents: {len(documents)}\n\n{joined}"
-    return retrieved
-###==========================================2-3
-
-def prepare_prompt(prompt, chat_history):
-    if genparam.TYPE == "chat" and chat_history:
-        chats = "\n".join([f"{message['role']}: \"{message['content']}\"" for message in chat_history])
-        prompt = f"""Retrieved Contextual Information:\n__grounding__\n\nConversation History:\n{chats}\n\nNew User Input: {prompt}"""
-        return prompt
+    if failure_matches >= 2:
+        return "The story ends in failure. Critical mistakes were made."
+    elif success_ratio >= 0.7:
+        return "The story concludes successfully! Well done!"
     else:
-        prompt = f"""Retrieved Contextual Information:\n__grounding__\n\nUser Input: {prompt}"""
-        return prompt
-
-def apply_prompt_syntax(prompt, system_prompt, prompt_template, bake_in_prompt_syntax):
-    model_family_syntax = {
-        "llama3-instruct (llama-3, 3.1 & 3.2) - system": """<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_prompt}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n""",
-        "llama3-instruct (llama-3, 3.1 & 3.2) - user": """<|begin_of_text|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n""",
-        "granite-13b-chat & instruct - system": """<|system|>\n{system_prompt}\n<|user|>\n{prompt}\n<|assistant|>\n\n""",
-        "granite-13b-chat & instruct - user": """<|user|>\n{prompt}\n<|assistant|>\n\n""",
-        "mistral & mixtral v2 tokenizer - system": """<s>[INST] System Prompt: {system_prompt} [/INST][INST] {prompt} [/INST]\n\n""",
-        "mistral & mixtral v2 tokenizer - user": """<s>[INST] {prompt} [/INST]\n\n""",
-        "no syntax - system": """{system_prompt}\n\n{prompt}""",
-        "no syntax - user": """{prompt}"""
-    }
-    
-    if bake_in_prompt_syntax:
-        template = model_family_syntax[prompt_template]
-        if system_prompt:
-            return template.format(system_prompt=system_prompt, prompt=prompt)
-    return prompt
-
-def generate_response(watsonx_llm, prompt_data, params):
-    generated_response = watsonx_llm.generate_text_stream(prompt=prompt_data, params=params)
-    for chunk in generated_response:
-        yield chunk
-
-def fetch_response(user_input, milvus_client, emb, vector_props, config, system_prompt, chat_history):
-    # Get grounding documents
-    grounding = proximity_search(
-        question=user_input,
-        milvus_client=milvus_client,
-        emb=emb,
-        vector_props=vector_promps,
-        config=config_path
-    )
-
-    # Special handling for PATH-er B. (first column)
-    if chat_history == st.session_state.chat_history_1:
-        # Display user question first
-        with st.chat_message("user", avatar=genparam.USER_AVATAR):
-            st.markdown(user_input)
-        
-        # Parse and display each document from the grounding
-        documents = grounding.split("\n\n")[2:]  # Skip the count line and first newline
-        for doc in documents:
-            if doc.strip():  # Only process non-empty strings
-                parts = doc.split("\n")
-                doc_name = parts[0].replace("Document: ", "")
-                content = parts[1].replace("Content: ", "")
-                
-                # Display document with delay
-                time.sleep(0.5)
-                st.markdown(f"**{doc_name}**")
-                st.code(content)
-
-        # Store in chat history
-        return grounding
-
-    # For MOD-ther S. (second column) and SYS-ter V. (third column)
-    else:
-        prompt = prepare_prompt(user_input, chat_history)
-        prompt_data = apply_prompt_syntax(
-            prompt,
-            system_prompt,  # Using the system_prompt passed to the function
-            get_active_prompt_template(),
-            genparam.BAKE_IN_PROMPT_SYNTAX
-        )
-        prompt_data = prompt_data.replace("__grounding__", grounding)
-
-        # # Add debug information to column 1 if enabled
-        # if genparam.INPUT_DEBUG_VIEW == 1:
-        #     with col1:  # Access first column
-        #         bot_name = genparam.BOT_2_NAME if chat_history == st.session_state.chat_history_2 else genparam.BOT_3_NAME
-        #         bot_avatar = genparam.BOT_2_AVATAR if chat_history == st.session_state.chat_history_2 else genparam.BOT_3_AVATAR
-        #         st.markdown(f"**{bot_avatar} {bot_name} Prompt Data:**")
-        #         st.code(prompt_data, language="text")
-
-    # Continue with normal processing for columns 2 and 3
-    watsonx_llm = ModelInference(
-        api_client=client, 
-        model_id=get_active_model(),
-        verify=genparam.VERIFY
-    )
+        return "The story ends with mixed results. Some opportunities were missed."
 
-    params = {
-        GenParams.DECODING_METHOD: genparam.DECODING_METHOD,
-        GenParams.MAX_NEW_TOKENS: genparam.MAX_NEW_TOKENS,
-        GenParams.MIN_NEW_TOKENS: genparam.MIN_NEW_TOKENS,
-        GenParams.REPETITION_PENALTY: genparam.REPETITION_PENALTY,
-        GenParams.STOP_SEQUENCES: genparam.STOP_SEQUENCES
-    }
-
-    bot_name = None
-    bot_avatar = None
-    if chat_history == st.session_state.chat_history_1:
-        bot_name = genparam.BOT_1_NAME
-        bot_avatar = genparam.BOT_1_AVATAR
-    elif chat_history == st.session_state.chat_history_2:
-        bot_name = genparam.BOT_2_NAME
-        bot_avatar = genparam.BOT_2_AVATAR
+def update_story_state(state: StoryState, user_input: str, char1_response: str, char2_response: str):
+    state.current_step += 1
+    state.user_inputs.append(user_input)
+    state.character1_responses.append(char1_response)
+    state.character2_responses.append(char2_response)
+    
+    # Check if story should end
+    if state.current_step >= state.max_steps:
+        state.story_outcome = evaluate_outcome(state, st.session_state.current_story)
+
+def generate_character_response(
+    character_name: str,
+    story_log: List[str],
+    user_input: str,
+    other_response: Optional[str] = None,
+    is_character1: bool = True
+) -> str:
+    # This would normally use watsonx.ai or another LLM
+    # For now, return placeholder responses
+    if is_character1:
+        return f"{character_name}: That's an interesting approach. Let's see where this leads..."
     else:
-        bot_name = genparam.BOT_3_NAME
-        bot_avatar = genparam.BOT_3_AVATAR
-
-    with st.chat_message(bot_name, avatar=bot_avatar):
-        if chat_history != st.session_state.chat_history_1:  # Only generate responses for columns 2 and 3
-            stream = generate_response(watsonx_llm, prompt_data, params)
-            response = st.write_stream(stream)
-            
-            # Only capture tokens for MOD-ther S. and SYS-ter V.
-            if genparam.TOKEN_CAPTURE_ENABLED and chat_history != st.session_state.chat_history_1:
-                token_stats = capture_tokens(prompt_data, response, bot_name)
-                if token_stats:
-                    st.session_state.token_statistics.append(token_stats)
-        else:
-            response = grounding  # For column 1, we already displayed the content
-
-    return response
-
-def capture_tokens(prompt_data, response, chat_number):
-    if not genparam.TOKEN_CAPTURE_ENABLED:
-        return
-
-    watsonx_llm = ModelInference(
-        api_client=client, 
-        model_id=genparam.SELECTED_MODEL,
-        verify=genparam.VERIFY
-    )
-
-    input_tokens = watsonx_llm.tokenize(prompt=prompt_data)["result"]["token_count"]
-    output_tokens = watsonx_llm.tokenize(prompt=response)["result"]["token_count"]
-    total_tokens = input_tokens + output_tokens
-
-    return {
-        "bot_name": bot_name,
-        "input_tokens": input_tokens,
-        "output_tokens": output_tokens,
-        "total_tokens": total_tokens,
-        "timestamp": time.strftime("%H:%M:%S")
-    }
+        return f"{character_name}: I have my doubts about this, but we'll see..."
 
 def main():
+    st.set_page_config(page_title="Interactive Story", layout="wide")
     initialize_session_state()
-
-    # Apply custom styles
-    st.markdown(three_column_style, unsafe_allow_html=True)
     
-    # Sidebar configuration
-    st.sidebar.header('The Solutioning Sages')
+    # Sidebar Configuration
+    st.sidebar.header('Story Selection')
     st.sidebar.divider()
     
-    # Knowledge Base Selection
-    selected_kb = st.sidebar.selectbox(
-        "Select Knowledge Base",
-        KNOWLEDGE_BASE_OPTIONS,
-        index=KNOWLEDGE_BASE_OPTIONS.index(st.session_state.selected_kb)
+    # Category Selection
+    selected_category = st.sidebar.selectbox(
+        "Select Story Category",
+        list(STORY_CATEGORIES.keys()),
+        index=list(STORY_CATEGORIES.keys()).index(st.session_state.selected_category)
     )
     
-    # Update knowledge base related values if selection changes
-    if selected_kb != st.session_state.selected_kb:
-        st.session_state.selected_kb = selected_kb
-        # Update the client with the new project_id
-        global client, wml_credentials
-        wml_credentials, client = setup_client(VECTOR_INDEXES[selected_kb]["project_id"])
+    # Update category and story if changed
+    if selected_category != st.session_state.selected_category:
+        st.session_state.selected_category = selected_category
+        st.session_state.current_story = random.choice(STORY_CATEGORIES[selected_category])
+        st.session_state.story_state = StoryState()  # Reset state for new story
+        st.rerun()
     
-    # Display current knowledge base contents
-    with st.sidebar.expander("Knowledge Base Contents"):
-        for doc in VECTOR_INDEXES[selected_kb]["contents"]:
-            st.write(f"📄 {doc}")
+    # Display available stories in category
+    with st.sidebar.expander(f"Available {selected_category} Stories"):
+        for story in STORY_CATEGORIES[selected_category]:
+            st.write(f"📖 {story['title']}")
     
-    # Display active model information
-    st.sidebar.divider()
-    active_model = genparam.SELECTED_MODEL_1 if genparam.ACTIVE_MODEL == 0 else genparam.SELECTED_MODEL_2
-    st.sidebar.markdown("**Active Model:**")
-    st.sidebar.code(active_model)
+    # Optional: Select specific story
+    specific_story = st.sidebar.selectbox(
+        "Select Specific Story",
+        [story["title"] for story in STORY_CATEGORIES[selected_category]],
+        index=[story["title"] for story in STORY_CATEGORIES[selected_category]].index(st.session_state.current_story["title"])
+    )
     
+    # Update if specific story changed
+    if specific_story != st.session_state.current_story["title"]:
+        st.session_state.current_story = STORY_LOOKUP[specific_story]
+        st.session_state.story_state = StoryState()  # Reset state for new story
+        st.rerun()
+    
+    # Display story stats
     st.sidebar.divider()
+    st.sidebar.subheader("Story Progress")
+    progress = (st.session_state.story_state.current_step / 5) * 100
+    st.sidebar.progress(progress)
+    st.sidebar.write(f"Step {st.session_state.story_state.current_step + 1}/5")
     
-    # Display token statistics in sidebar
-    st.sidebar.subheader("Token Usage Statistics")
+    # Create three columns
+    col1, col2, col3 = st.columns(3)
     
-    # Group token statistics by interaction (for MOD-ther S. and SYS-ter V. only)
-    if st.session_state.token_statistics:
-        current_timestamp = None
-        interaction_count = 0
-        stats_by_time = {}
+    # Story Progress Column
+    with col1:
+        st.header("Story Progress")
+        st.write(f"**{st.session_state.current_story['title']}**")
+        st.write(st.session_state.current_story['initial_setup'])
         
-        # Group stats by timestamp
-        for stat in st.session_state.token_statistics:
-            if stat["timestamp"] not in stats_by_time:
-                stats_by_time[stat["timestamp"]] = []
-            stats_by_time[stat["timestamp"]].append(stat)
+        # Display story log
+        for step_num, (step, user_input) in enumerate(zip(
+            st.session_state.current_story['steps'][:st.session_state.story_state.current_step],
+            st.session_state.story_state.user_inputs
+        )):
+            st.write(f"Step {step_num + 1}: {step}")
+            st.write(f"Your action: {user_input}")
+            st.write("---")
         
-        # Display grouped stats
-        for timestamp, stats in stats_by_time.items():
-            interaction_count += 1
-            st.sidebar.markdown(f"**Interaction {interaction_count}** ({timestamp})")
-            
-            # Calculate total tokens for this interaction
-            total_input = sum(stat['input_tokens'] for stat in stats)
-            total_output = sum(stat['output_tokens'] for stat in stats)
-            total = total_input + total_output
-            
-            # Display individual bot statistics
-            for stat in stats:
-                st.sidebar.markdown(
-                    f"_{stat['bot_name']}_  \n"
-                    f"Input: {stat['input_tokens']} tokens  \n"
-                    f"Output: {stat['output_tokens']} tokens  \n"
-                    f"Total: {stat['total_tokens']} tokens"
-                )
-            
-            # Display interaction totals
-            st.sidebar.markdown("**Interaction Totals:**")
-            st.sidebar.markdown(
-                f"Total Input: {total_input} tokens  \n"
-                f"Total Output: {total_output} tokens  \n"
-                f"Total Usage: {total} tokens"
-            )
-            st.sidebar.markdown("---")
+        # Display current step if story isn't finished
+        if st.session_state.story_state.story_outcome is None:
+            current_step = st.session_state.current_story['steps'][st.session_state.story_state.current_step]
+            st.write(f"Current Situation: {current_step}")
     
-    st.sidebar.markdown("")
-
-    if not check_password():
-        st.stop()
-
-    # Get user input before column creation
-    user_input = st.chat_input("Ask your question here", key="user_input")
+    # Character 1 Column
+    with col2:
+        st.header(st.session_state.current_story['character1_name'])
+        for response in st.session_state.story_state.character1_responses:
+            st.write(response)
+    
+    # Character 2 Column
+    with col3:
+        st.header(st.session_state.current_story['character2_name'])
+        for response in st.session_state.story_state.character2_responses:
+            st.write(response)
     
-    if user_input:
-        # Create three columns
-        col1, col2, col3 = st.columns(3)
+    # User Input Section
+    if st.session_state.story_state.story_outcome is None:
+        user_input = st.text_input(
+            "What do you do?",
+            key=f"user_input_{st.session_state.story_state.current_step}"
+        )
         
-        # First column - PATH-er B. (Document Display)
-        with col1:
-            st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
-            st.subheader(f"{genparam.BOT_1_AVATAR} {genparam.BOT_1_NAME}")
-            st.markdown("<div class='chat-messages'>", unsafe_allow_html=True)
-            
-            # Display previous messages
-            for message in st.session_state.chat_history_1:
-                if message["role"] == "user":
-                    with st.chat_message(message["role"], avatar=genparam.USER_AVATAR):
-                        st.markdown(message['content'])
-                else:
-                    # Parse and display stored documents
-                    documents = message['content'].split("\n\n")[2:]  # Skip count line
-                    for doc in documents:
-                        if doc.strip():
-                            parts = doc.split("\n")
-                            doc_name = parts[0].replace("Document: ", "")
-                            content = parts[1].replace("Content: ", "")
-                            st.markdown(f"**{doc_name}**")
-                            st.code(content)
-            
-            # Add user message and get new response
-            st.session_state.chat_history_1.append({"role": "user", "content": user_input, "avatar": genparam.USER_AVATAR})
-            milvus_client, emb, vector_index_properties, vector_store_schema = setup_vector_index(
-                client, 
-                wml_credentials,
-                VECTOR_INDEXES[st.session_state.selected_kb]["index_1"]
+        if user_input:
+            # Generate character responses
+            char1_response = generate_character_response(
+                st.session_state.current_story['character1_name'],
+                st.session_state.story_state.story_log,
+                user_input
             )
-            system_prompt = genparam.BOT_1_PROMPT
             
-            response = fetch_response(
-                user_input, 
-                milvus_client, 
-                emb, 
-                vector_index_properties, 
-                vector_store_schema,
-                system_prompt,
-                st.session_state.chat_history_1
+            char2_response = generate_character_response(
+                st.session_state.current_story['character2_name'],
+                st.session_state.story_state.story_log,
+                user_input,
+                char1_response,
+                False
             )
-            st.session_state.chat_history_1.append({"role": genparam.BOT_1_NAME, "content": response, "avatar": genparam.BOT_1_AVATAR})
-            st.markdown("</div></div>", unsafe_allow_html=True)
-
-        # Second column - MOD-ther S. (Uses documents from first vector index)
-        with col2:
-            st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
-            st.subheader(f"{genparam.BOT_2_AVATAR} {genparam.BOT_2_NAME}")
-            st.markdown("<div class='chat-messages'>", unsafe_allow_html=True)
-            
-            for message in st.session_state.chat_history_2:
-                if message["role"] != "user":
-                    with st.chat_message(message["role"], avatar=genparam.BOT_2_AVATAR):
-                        st.markdown(message['content'])
-            
-            st.session_state.chat_history_2.append({"role": "user", "content": user_input, "avatar": genparam.USER_AVATAR})
-            milvus_client, emb, vector_index_properties, vector_store_schema = setup_vector_index(
-                client, 
-                wml_credentials,
-                #VECTOR_INDEXES[st.session_state.selected_kb]["index_1"]
-                config_path
-            )
-            system_prompt = SYSTEM_PROMPTS[st.session_state.selected_kb]["bot_2"]
-            
-            response = fetch_response(
-                user_input, 
-                milvus_client, 
-                emb, 
-                vector_index_properties, 
-                vector_store_schema,
-                system_prompt,
-                st.session_state.chat_history_2
-            )
-
-            if genparam.INPUT_DEBUG_VIEW == 1:
-                with col1:  # Access first column
-                    bot_name = genparam.BOT_2_NAME if st.session_state.chat_history_1 == st.session_state.chat_history_2 else genparam.BOT_3_NAME
-                    bot_avatar = genparam.BOT_2_AVATAR if st.session_state.chat_history_1 == st.session_state.chat_history_2 else genparam.BOT_3_AVATAR
-                    st.markdown(f"**{bot_avatar} {bot_name} Prompt Data:**")
-                    st.code(prompt_data, language="text")
-            
-            st.session_state.chat_history_2.append({"role": genparam.BOT_2_NAME, "content": response, "avatar": genparam.BOT_2_AVATAR})
-            st.markdown("</div></div>", unsafe_allow_html=True)
-            
-        # Third column - SYS-ter V. (Uses second vector index and chat history from second column)
-        with col3:
-            st.markdown("<div class='chat-container'>", unsafe_allow_html=True)
-            st.subheader(f"{genparam.BOT_3_AVATAR} {genparam.BOT_3_NAME}")
-            st.markdown("<div class='chat-messages'>", unsafe_allow_html=True)
             
-            for message in st.session_state.chat_history_3:
-                if message["role"] != "user":
-                    with st.chat_message(message["role"], avatar=genparam.BOT_3_AVATAR):
-                        st.markdown(message['content'])
-            
-            st.session_state.chat_history_3.append({"role": "user", "content": user_input, "avatar": genparam.USER_AVATAR})
-            milvus_client, emb, vector_index_properties, vector_store_schema = setup_vector_index(
-                client, 
-                wml_credentials,
-                VECTOR_INDEXES[st.session_state.selected_kb]["index_2"]
-            )
-            system_prompt = SYSTEM_PROMPTS[st.session_state.selected_kb]["bot_3"]
-            
-            response = fetch_response(
-                user_input, 
-                milvus_client, 
-                emb, 
-                vector_index_properties, 
-                vector_store_schema,
-                system_prompt,
-                st.session_state.chat_history_3
-            )
-            
-            if genparam.INPUT_DEBUG_VIEW == 1:
-                with col1:  # Access first column
-                    bot_name = genparam.BOT_2_NAME if st.session_state.chat_history_1 == st.session_state.chat_history_2 else genparam.BOT_3_NAME
-                    bot_avatar = genparam.BOT_2_AVATAR if st.session_state.chat_history_1 == st.session_state.chat_history_2 else genparam.BOT_3_AVATAR
-                    st.markdown(f"**{bot_avatar} {bot_name} Prompt Data:**")
-                    st.code(prompt_data, language="text")
-            
-            st.session_state.chat_history_3.append({"role": genparam.BOT_3_NAME, "content": response, "avatar": genparam.BOT_3_AVATAR})
-            st.markdown("</div></div>", unsafe_allow_html=True)
-
-        # Update sidebar with new question
-        st.sidebar.markdown("---")
-        st.sidebar.markdown("**Latest Question:**")
-        st.sidebar.markdown(f"_{user_input}_")
+            # Update state
+            update_story_state(st.session_state.story_state, user_input, char1_response, char2_response)
+            st.rerun()
+    else:
+        # Display story outcome
+        st.write(st.session_state.story_state.story_outcome)
+        if st.button("Start New Story"):
+            st.session_state.story_state = StoryState()
+            st.session_state.current_story = random.choice(MICRO_STORIES)
+            st.rerun()
 
 if __name__ == "__main__":
     main()