diff --git a/ai_agent_tutorials/ai_r1-tooluse-langroid/main.py b/ai_agent_tutorials/ai_r1-tooluse-langroid/main.py index 0b6c599..e94023f 100644 --- a/ai_agent_tutorials/ai_r1-tooluse-langroid/main.py +++ b/ai_agent_tutorials/ai_r1-tooluse-langroid/main.py @@ -18,17 +18,7 @@ CLAUDE_MODEL: str = "claude-3-5-sonnet-20241022" load_dotenv() class ModelChain: - """ - A class to handle interactions with DeepSeek and Claude models. - """ def __init__(self, deepseek_api_key: str, anthropic_api_key: str) -> None: - """ - Initialize the ModelChain with API keys. - - Args: - deepseek_api_key: API key for DeepSeek - anthropic_api_key: API key for Anthropic - """ self.deepseek_client = OpenAI( api_key=deepseek_api_key, base_url="https://api.deepseek.com" @@ -47,21 +37,9 @@ class ModelChain: return self.current_model def get_deepseek_reasoning(self, user_input: str) -> str: - """ - Get reasoning from DeepSeek model. - - Args: - user_input: User's input text - - Returns: - str: Reasoning content from DeepSeek - """ start_time = time.time() self.deepseek_messages.append({"role": "user", "content": user_input}) - if self.show_reasoning: - rprint("\n[blue]Reasoning Process[/]") - response = self.deepseek_client.chat.completions.create( model=DEEPSEEK_MODEL, max_tokens=1, @@ -71,27 +49,23 @@ class ModelChain: reasoning_content = "" final_content = "" - reasoning_placeholder = st.empty() - for chunk in response: - if chunk.choices[0].delta.reasoning_content: - reasoning_piece = chunk.choices[0].delta.reasoning_content - reasoning_content += reasoning_piece - reasoning_placeholder.markdown(reasoning_content) - if self.show_reasoning: - print(reasoning_piece, end="", flush=True) - elif chunk.choices[0].delta.content: - final_content += chunk.choices[0].delta.content + # Create expander for reasoning + with st.expander("💭 Reasoning Process", expanded=True): + reasoning_placeholder = st.empty() + + for chunk in response: + if chunk.choices[0].delta.reasoning_content: + reasoning_piece = chunk.choices[0].delta.reasoning_content + reasoning_content += reasoning_piece + reasoning_placeholder.markdown(reasoning_content) + elif chunk.choices[0].delta.content: + final_content += chunk.choices[0].delta.content - elapsed_time = time.time() - start_time - if elapsed_time >= 60: - time_str = f"{elapsed_time/60:.1f} minutes" - else: - time_str = f"{elapsed_time:.1f} seconds" - rprint(f"\n\n[yellow]Thought for {time_str}[/]") + elapsed_time = time.time() - start_time + time_str = f"{elapsed_time/60:.1f} minutes" if elapsed_time >= 60 else f"{elapsed_time:.1f} seconds" + st.caption(f"⏱️ Thought for {time_str}") - if self.show_reasoning: - print("\n") return reasoning_content def get_claude_response(self, user_input: str, reasoning: str) -> str: @@ -116,46 +90,46 @@ class ModelChain: } messages = [user_message, assistant_prefill] - response_placeholder = st.empty() - rprint(f"[green]{self.get_model_display_name()}[/]", end="") - try: - with self.claude_client.messages.stream( - model=self.current_model, - messages=messages, - max_tokens=8000 - ) as stream: - full_response = "" - for text in stream.text_stream: - full_response += text - response_placeholder.markdown(full_response) + # Create expander for Claude's response + with st.expander("🤖 Claude's Response", expanded=True): + response_placeholder = st.empty() + + with self.claude_client.messages.stream( + model=self.current_model, + messages=messages, + max_tokens=8000 + ) as stream: + full_response = "" + for text in stream.text_stream: + full_response += text + response_placeholder.markdown(full_response) - self.claude_messages.extend([user_message, { - "role": "assistant", - "content": [{"type": "text", "text": full_response}] - }]) - self.deepseek_messages.append({"role": "assistant", "content": full_response}) + self.claude_messages.extend([user_message, { + "role": "assistant", + "content": [{"type": "text", "text": full_response}] + }]) + self.deepseek_messages.append({"role": "assistant", "content": full_response}) - print("\n") - return full_response + return full_response except Exception as e: - rprint(f"\n[red]Error in response: {str(e)}[/]") + st.error(f"Error in response: {str(e)}") return "Error occurred while getting response" def main() -> None: """Main function to run the Streamlit app.""" - st.title("AI Assistant") + st.title("🤖 AI Assistant") # Sidebar for API keys with st.sidebar: - st.header("API Configuration") + st.header("⚙️ Configuration") deepseek_api_key = st.text_input("DeepSeek API Key", type="password") anthropic_api_key = st.text_input("Anthropic API Key", type="password") show_reasoning = st.toggle("Show Reasoning Process", value=True) - if st.button("Clear Chat History"): + if st.button("🗑️ Clear Chat History"): st.session_state.messages = [] st.experimental_rerun() @@ -171,7 +145,7 @@ def main() -> None: # Chat input if prompt := st.chat_input("What would you like to know?"): if not deepseek_api_key or not anthropic_api_key: - st.error("Please enter both API keys in the sidebar.") + st.error("⚠️ Please enter both API keys in the sidebar.") return # Initialize ModelChain @@ -185,12 +159,14 @@ def main() -> None: # Get AI response with st.chat_message("assistant"): if show_reasoning: - reasoning = chain.get_deepseek_reasoning(prompt) + with st.spinner("🤔 Thinking..."): + reasoning = chain.get_deepseek_reasoning(prompt) else: reasoning = "" - response = chain.get_claude_response(prompt, reasoning) - st.session_state.messages.append({"role": "assistant", "content": response}) + with st.spinner("✍️ Responding..."): + response = chain.get_claude_response(prompt, reasoning) + st.session_state.messages.append({"role": "assistant", "content": response}) if __name__ == "__main__": main() \ No newline at end of file