third push

.env_example

@@ -1,2 +1,2 @@
 MODEL_NAME=openai/gpt-3.5-turbo
-MODEL_API_KEY=sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+MODEL_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx

app/app.py

@@ -2,11 +2,12 @@
 from calendar import c
 from dataclasses import dataclass
 from math import exp
+from shutil import which
+import time
 from webbrowser import get
 from litellm.types.utils import ModelResponse
 import streamlit as st
-from streamlit.runtime.uploaded_file_manager import UploadedFile
-from app.utils import generate_answer, dict_to_markdown
+from app.utils import generate_answer, load_llm
 from core.types import ThoughtStepsDisplay, ThoughtSteps, BigMessage , Message
 from .app_config import InputConfig, ENV_FILE_PATH, CONFIG_FILE_PATH
 from core.llms.base_llm import BaseLLM
@@ -16,18 +17,18 @@ from PIL import Image
 from core.prompts.think_mark_think import SYSTEM_PROMPT
 
 st.set_page_config(page_title="Open-o1", page_icon="🧠", layout="wide")
-st.title('Open-O1')
-st.write('Welcome to Open-O1!')
+
 
 def config_sidebar(config:InputConfig) -> InputConfig:
     st.sidebar.header('Configuration')
-    model_name =    st.sidebar.text_input('Enter Model Name: e.g. provider/model-name',value=config.model_name, placeholder='openai/gpt-3.5-turbo')
-    model_api_key = st.sidebar.text_input('Enter API Key: ',type='password',value=config.model_api_key, placeholder='sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx')
-    max_tokens =    st.sidebar.number_input('Enter Max Tokens per Thought: ',value=config.max_tokens, min_value=1)
-    max_steps =     st.sidebar.number_input('Enter Max Thinking Steps: ',value=config.max_steps, min_value=1, step=1, ) 
-    temperature =   st.sidebar.number_input('Enter Temperature: ',value=config.temperature, min_value=0.0, step=0.1, max_value=10.0)
-    timeout =       st.sidebar.number_input('Enter timeout(seconds): ',value=config.timeout, min_value=0.0,step = 1.0)
-    sleeptime =     st.sidebar.number_input('Enter Sleep Time(seconds): (time bw requests to avoid rate limit)',value=config.sleeptime, min_value=0.0, step = 1.0) 
+    model_name =    st.sidebar.text_input('Model Name: e.g. provider/model-name',value=config.model_name, placeholder='openai/gpt-3.5-turbo')
+    model_api_key = st.sidebar.text_input('API Key: ',type='password',value=config.model_api_key, placeholder='sk-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx')
+    max_tokens =    st.sidebar.number_input('Max Tokens per Thought: ',value=config.max_tokens, min_value=1)
+    max_steps =     st.sidebar.number_input('Max Thinking Steps: ',value=config.max_steps, min_value=1, step=1, ) 
+    temperature =   st.sidebar.number_input('Temperature: ',value=config.temperature, min_value=0.0, step=0.1, max_value=10.0)
+    timeout =       st.sidebar.number_input('Timeout(seconds): ',value=config.timeout, min_value=0.0,step = 1.0)
+    sleeptime =     st.sidebar.number_input('Sleep Time(seconds)',value=config.sleeptime, min_value=0.0, step = 1.0, help='Time between requests to avoid hitting rate limit')  
+    force_max_steps = st.sidebar.checkbox('Force Max Steps', value=config.force_max_steps, help="If checked, will generate given number of max steps. If not checked, assistant can stop at few step thinking it has the write answer.") 
     
     config.model_name = model_name
     config.model_api_key = model_api_key
@@ -36,6 +37,8 @@ def config_sidebar(config:InputConfig) -> InputConfig:
     config.temperature = temperature
     config.timeout = timeout
     config.sleeptime = sleeptime
+    config.force_max_steps = force_max_steps
+    
 
     if st.sidebar.button('Save config'):
         config.save(env_file=ENV_FILE_PATH, config_file=CONFIG_FILE_PATH)
@@ -43,43 +46,27 @@ def config_sidebar(config:InputConfig) -> InputConfig:
         
     return config
 
-def load_llm(config:InputConfig, tools=None) -> BaseLLM:
-    return LLM(api_key=config.model_api_key, model=config.model_name, tools=tools)
-
-    
-def image_buffer_to_pillow_image(image_buffer:UploadedFile) -> Image.Image:
-    return Image.open(image_buffer)
-
-
-def process_user_input(user_input:str, image:Image.Image=None)->dict:
-    if image:
-        message = [user_message_with_images(user_msg_str=user_input, images=[image])] 
-    else:
-        message = [{"role": "user", "content": user_input}]
-    return message
-
-
 
     
 def main():
+    st.title('Open-O1')
+    st.write('Welcome to Open-O1!')
+
     
     config = InputConfig.load(env_file=ENV_FILE_PATH, config_file=CONFIG_FILE_PATH)    
     config = config_sidebar(config=config)
     llm = load_llm(config)
 
-    current_tab = ''
     
     current_tab='o1_tab'
-    messages_attr_name = f"{current_tab}_messages"
     big_message_attr_name = f"{current_tab}_big_messages"
     
 
     clear_chat_bt = st.sidebar.button('Clear Chat')
     if clear_chat_bt:
-        delattr(st.session_state, messages_attr_name)
+        delattr(st.session_state, big_message_attr_name)
 
 
-    message_attr = set_and_get_state_attr(messages_attr_name, default_value=[])
     big_message_attr = set_and_get_state_attr(big_message_attr_name, default_value=[])
     
     # this prints the older messages
@@ -93,7 +80,7 @@ def main():
                 if message.role == 'user':
                     st.markdown(message.content)
                 else:
-                    print_thought(message.to_thought_steps_display(), is_final=True)
+                    print_thought(message.content.to_thought_steps_display(), is_final=True)
 
             
     
@@ -118,30 +105,42 @@ def main():
                 if message["role"] == "user":
                     message["content"] = f"{message['content']}, json format"
             
+            start_time = time.time()
             
-            for num,step in enumerate(generate_answer(
-                messages=messages, 
-                max_steps=config.max_steps, 
-                stream=False, 
-                max_tokens=config.max_tokens, 
-                temperature=config.temperature, 
-                sleeptime=config.sleeptime,
-                timeout=config.timeout, 
-                llm=llm,
-                response_format={ "type": "json_object" }
-                
-                ),1):
-
-                thoughts.append(step)
-                print_thought(step.to_thought_steps_display(), is_final=False)
+            with st.status("Thinking...", expanded=True) as status:
+
+                for step in generate_answer(
+                    messages=messages, 
+                    max_steps=config.max_steps, 
+                    stream=False, 
+                    max_tokens=config.max_tokens, 
+                    temperature=config.temperature, 
+                    sleeptime=config.sleeptime,
+                    timeout=config.timeout, 
+                    llm=llm,
+                    force_max_steps=config.force_max_steps,
+                    response_format={ "type": "json_object" }
+                    
+                    ):
+
+                    thoughts.append(step)
+
+                    st.write(step.to_thought_steps_display().md())
+                    
+                    status.update(label=step.step_title, state="running", expanded=False)
+
+                status.update(
+                    label=f"Thought for {time.time()-start_time:.2f} seconds", state="complete", expanded=False
+                )
 
             last_step = thoughts.pop()
-            message_attr.append(BigMessage(
+            print_thought(last_step.to_thought_steps_display(), is_final=True)
+
+            big_message_attr.append(BigMessage(
                 role="assistant", 
                 content=last_step, 
                 thoughts=thoughts
             ))
-            # st.markdown(dict_to_markdown(step.model_dump()))
 
 
     
@@ -155,8 +154,9 @@ def print_thought(thought:ThoughtStepsDisplay, is_final:bool=False):
     if is_final:
        st.markdown(thought.md())
     else:
-        st.markdown(f'\n```json\n{thought.model_dump_json()}\n```\n', unsafe_allow_html=True) 
-        
+        # st.markdown(f'\n```json\n{thought.model_dump_json()}\n```\n', unsafe_allow_html=True) 
+        with st.expander(f'{thought.step_title}'):
+            st.markdown(thought.md())
 
         
     

app/app_config.py

@@ -24,6 +24,7 @@ class InputConfig:
     temperature: float = 0.2
     timeout: float = 30.0
     sleeptime: float = 0.0
+    force_max_steps: bool = True
 
     @classmethod
     def load(cls, env_file=ENV_FILE_PATH, config_file=CONFIG_FILE_PATH):
@@ -46,7 +47,8 @@ class InputConfig:
             max_steps=config_dict.get('max_steps', cls.max_steps),
             temperature=config_dict.get('temperature', cls.temperature),
             timeout=config_dict.get('timeout', cls.timeout),
-            sleeptime=config_dict.get('sleeptime', cls.sleeptime)
+            sleeptime=config_dict.get('sleeptime', cls.sleeptime),
+            force_max_steps=config_dict.get('force_max_steps', cls.force_max_steps)
         )
 
     def save(self, env_file=ENV_FILE_PATH, config_file=CONFIG_FILE_PATH):
@@ -74,7 +76,8 @@ class InputConfig:
             'max_steps': self.max_steps,
             'temperature': self.temperature,
             'timeout': self.timeout,
-            'sleeptime': self.sleeptime
+            'sleeptime': self.sleeptime,
+            'force_max_steps': self.force_max_steps
         }
         with open(config_file, 'w') as f:
             json.dump(config_dict, f, indent=4)

app/utils.py

@@ -11,10 +11,16 @@ import os
 import time
 from core.utils import parse_with_fallback
 from termcolor import colored
+from app.app_config import InputConfig
+from core.llms.litellm_llm import LLM
+from core.llms.utils import user_message_with_images
+from PIL import Image
+from streamlit.runtime.uploaded_file_manager import UploadedFile
 
 
 
-def generate_answer(messages: list[dict], max_steps: int = 20, llm: BaseLLM = None, sleeptime: float = 0.0, **kwargs):
+
+def generate_answer(messages: list[dict], max_steps: int = 20, llm: BaseLLM = None, sleeptime: float = 0.0, force_max_steps: bool = False, **kwargs):
     thoughts = []
     
     for i in range(max_steps):
@@ -24,13 +30,16 @@ def generate_answer(messages: list[dict], max_steps: int = 20, llm: BaseLLM = No
         
         print(colored(f"{i+1} - {response}", 'yellow'))
 
-        if thought:
-            thoughts.append(thought)
-            messages.append({"role": "assistant", "content": thought.model_dump_json()})
-            messages.append({"role": "user", "content": REVIEW_PROMPT})
+        thoughts.append(thought)
+        messages.append({"role": "assistant", "content": thought.model_dump_json()})
+        yield thought
+                
+        if thought.is_final_answer and not thought.next_step and not force_max_steps:
+            break
+        
+        messages.append({"role": "user", "content": REVIEW_PROMPT})
 
-            yield thought
-            time.sleep(sleeptime)
+        time.sleep(sleeptime)
 
     # Get the final answer after all thoughts are processed
     messages += [{"role": "user", "content": FINAL_ANSWER_PROMPT}]
@@ -66,3 +75,22 @@ def dict_to_markdown(d:dict) -> str:
         md += f"{value}\n"
     return md
 
+
+
+
+def load_llm(config:InputConfig, tools=None) -> BaseLLM:
+    return LLM(api_key=config.model_api_key, model=config.model_name, tools=tools)
+
+    
+def image_buffer_to_pillow_image(image_buffer:UploadedFile) -> Image.Image:
+    return Image.open(image_buffer)
+
+
+def process_user_input(user_input:str, image:Image.Image=None)->dict:
+    if image:
+        message = [user_message_with_images(user_msg_str=user_input, images=[image])] 
+    else:
+        message = [{"role": "user", "content": user_input}]
+    return message
+
+

core/types.py

@@ -30,15 +30,15 @@ class ThoughtStepsDisplay(BaseModel):
 
     
     def md(self):
-        return dedent(f'''
-            {self.step_title}
-            ### Thought
+        return f'''
+            #### {self.step_title}
+            
             {self.thought}
-            ### Answer
+            
             {self.answer}
-            ### Critic
+            
             {self.critic}
-            ''')
+            '''
 
 
 
@@ -56,3 +56,21 @@ class BigMessage(BaseModel):
 class Message(BaseModel):
     role:str
     content:str
+
+class InputConfig(BaseModel):
+    prompt: str = Field(..., description="prompt to use")
+    model: str = Field(..., description="model to use")
+    max_tokens: int = Field(..., description="max tokens to use")
+    temperature: float = Field(..., description="temperature to use")
+    top_p: float = Field(..., description="top p to use")
+    n: int = Field(..., description="number of responses to generate")
+    stream: bool = Field(..., description="whether to stream the response")
+    stop: list[str] | None = Field(..., description="stop sequences")
+    force_max_steps: bool = Field(False, description="force max steps")
+
+    def to_dict(self):
+        return self.model_dump()
+
+    def to_json(self):
+        return json.dumps(self.to_dict())
+

input_config.json

@@ -3,5 +3,6 @@
     "max_steps": 3,
     "temperature": 0.2,
     "timeout": 30.0,
-    "sleeptime": 2.0
+    "sleeptime": 2.0,
+    "force_max_steps": true
 }