feature 1 - requirmements.txt
This commit is contained in:
parent
4a69ac7224
commit
b87c6539cb
1 changed files with 43 additions and 67 deletions
|
|
@ -18,17 +18,7 @@ CLAUDE_MODEL: str = "claude-3-5-sonnet-20241022"
|
||||||
load_dotenv()
|
load_dotenv()
|
||||||
|
|
||||||
class ModelChain:
|
class ModelChain:
|
||||||
"""
|
|
||||||
A class to handle interactions with DeepSeek and Claude models.
|
|
||||||
"""
|
|
||||||
def __init__(self, deepseek_api_key: str, anthropic_api_key: str) -> None:
|
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(
|
self.deepseek_client = OpenAI(
|
||||||
api_key=deepseek_api_key,
|
api_key=deepseek_api_key,
|
||||||
base_url="https://api.deepseek.com"
|
base_url="https://api.deepseek.com"
|
||||||
|
|
@ -47,21 +37,9 @@ class ModelChain:
|
||||||
return self.current_model
|
return self.current_model
|
||||||
|
|
||||||
def get_deepseek_reasoning(self, user_input: str) -> str:
|
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()
|
start_time = time.time()
|
||||||
self.deepseek_messages.append({"role": "user", "content": user_input})
|
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(
|
response = self.deepseek_client.chat.completions.create(
|
||||||
model=DEEPSEEK_MODEL,
|
model=DEEPSEEK_MODEL,
|
||||||
max_tokens=1,
|
max_tokens=1,
|
||||||
|
|
@ -71,27 +49,23 @@ class ModelChain:
|
||||||
|
|
||||||
reasoning_content = ""
|
reasoning_content = ""
|
||||||
final_content = ""
|
final_content = ""
|
||||||
reasoning_placeholder = st.empty()
|
|
||||||
|
|
||||||
for chunk in response:
|
# Create expander for reasoning
|
||||||
if chunk.choices[0].delta.reasoning_content:
|
with st.expander("💭 Reasoning Process", expanded=True):
|
||||||
reasoning_piece = chunk.choices[0].delta.reasoning_content
|
reasoning_placeholder = st.empty()
|
||||||
reasoning_content += reasoning_piece
|
|
||||||
reasoning_placeholder.markdown(reasoning_content)
|
for chunk in response:
|
||||||
if self.show_reasoning:
|
if chunk.choices[0].delta.reasoning_content:
|
||||||
print(reasoning_piece, end="", flush=True)
|
reasoning_piece = chunk.choices[0].delta.reasoning_content
|
||||||
elif chunk.choices[0].delta.content:
|
reasoning_content += reasoning_piece
|
||||||
final_content += chunk.choices[0].delta.content
|
reasoning_placeholder.markdown(reasoning_content)
|
||||||
|
elif chunk.choices[0].delta.content:
|
||||||
|
final_content += chunk.choices[0].delta.content
|
||||||
|
|
||||||
elapsed_time = time.time() - start_time
|
elapsed_time = time.time() - start_time
|
||||||
if elapsed_time >= 60:
|
time_str = f"{elapsed_time/60:.1f} minutes" if elapsed_time >= 60 else f"{elapsed_time:.1f} seconds"
|
||||||
time_str = f"{elapsed_time/60:.1f} minutes"
|
st.caption(f"⏱️ Thought for {time_str}")
|
||||||
else:
|
|
||||||
time_str = f"{elapsed_time:.1f} seconds"
|
|
||||||
rprint(f"\n\n[yellow]Thought for {time_str}[/]")
|
|
||||||
|
|
||||||
if self.show_reasoning:
|
|
||||||
print("\n")
|
|
||||||
return reasoning_content
|
return reasoning_content
|
||||||
|
|
||||||
def get_claude_response(self, user_input: str, reasoning: str) -> str:
|
def get_claude_response(self, user_input: str, reasoning: str) -> str:
|
||||||
|
|
@ -116,46 +90,46 @@ class ModelChain:
|
||||||
}
|
}
|
||||||
|
|
||||||
messages = [user_message, assistant_prefill]
|
messages = [user_message, assistant_prefill]
|
||||||
response_placeholder = st.empty()
|
|
||||||
|
|
||||||
rprint(f"[green]{self.get_model_display_name()}[/]", end="")
|
|
||||||
|
|
||||||
try:
|
try:
|
||||||
with self.claude_client.messages.stream(
|
# Create expander for Claude's response
|
||||||
model=self.current_model,
|
with st.expander("🤖 Claude's Response", expanded=True):
|
||||||
messages=messages,
|
response_placeholder = st.empty()
|
||||||
max_tokens=8000
|
|
||||||
) as stream:
|
with self.claude_client.messages.stream(
|
||||||
full_response = ""
|
model=self.current_model,
|
||||||
for text in stream.text_stream:
|
messages=messages,
|
||||||
full_response += text
|
max_tokens=8000
|
||||||
response_placeholder.markdown(full_response)
|
) as stream:
|
||||||
|
full_response = ""
|
||||||
|
for text in stream.text_stream:
|
||||||
|
full_response += text
|
||||||
|
response_placeholder.markdown(full_response)
|
||||||
|
|
||||||
self.claude_messages.extend([user_message, {
|
self.claude_messages.extend([user_message, {
|
||||||
"role": "assistant",
|
"role": "assistant",
|
||||||
"content": [{"type": "text", "text": full_response}]
|
"content": [{"type": "text", "text": full_response}]
|
||||||
}])
|
}])
|
||||||
self.deepseek_messages.append({"role": "assistant", "content": full_response})
|
self.deepseek_messages.append({"role": "assistant", "content": full_response})
|
||||||
|
|
||||||
print("\n")
|
return full_response
|
||||||
return full_response
|
|
||||||
|
|
||||||
except Exception as e:
|
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"
|
return "Error occurred while getting response"
|
||||||
|
|
||||||
def main() -> None:
|
def main() -> None:
|
||||||
"""Main function to run the Streamlit app."""
|
"""Main function to run the Streamlit app."""
|
||||||
st.title("AI Assistant")
|
st.title("🤖 AI Assistant")
|
||||||
|
|
||||||
# Sidebar for API keys
|
# Sidebar for API keys
|
||||||
with st.sidebar:
|
with st.sidebar:
|
||||||
st.header("API Configuration")
|
st.header("⚙️ Configuration")
|
||||||
deepseek_api_key = st.text_input("DeepSeek API Key", type="password")
|
deepseek_api_key = st.text_input("DeepSeek API Key", type="password")
|
||||||
anthropic_api_key = st.text_input("Anthropic API Key", type="password")
|
anthropic_api_key = st.text_input("Anthropic API Key", type="password")
|
||||||
show_reasoning = st.toggle("Show Reasoning Process", value=True)
|
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.session_state.messages = []
|
||||||
st.experimental_rerun()
|
st.experimental_rerun()
|
||||||
|
|
||||||
|
|
@ -171,7 +145,7 @@ def main() -> None:
|
||||||
# Chat input
|
# Chat input
|
||||||
if prompt := st.chat_input("What would you like to know?"):
|
if prompt := st.chat_input("What would you like to know?"):
|
||||||
if not deepseek_api_key or not anthropic_api_key:
|
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
|
return
|
||||||
|
|
||||||
# Initialize ModelChain
|
# Initialize ModelChain
|
||||||
|
|
@ -185,12 +159,14 @@ def main() -> None:
|
||||||
# Get AI response
|
# Get AI response
|
||||||
with st.chat_message("assistant"):
|
with st.chat_message("assistant"):
|
||||||
if show_reasoning:
|
if show_reasoning:
|
||||||
reasoning = chain.get_deepseek_reasoning(prompt)
|
with st.spinner("🤔 Thinking..."):
|
||||||
|
reasoning = chain.get_deepseek_reasoning(prompt)
|
||||||
else:
|
else:
|
||||||
reasoning = ""
|
reasoning = ""
|
||||||
|
|
||||||
response = chain.get_claude_response(prompt, reasoning)
|
with st.spinner("✍️ Responding..."):
|
||||||
st.session_state.messages.append({"role": "assistant", "content": response})
|
response = chain.get_claude_response(prompt, reasoning)
|
||||||
|
st.session_state.messages.append({"role": "assistant", "content": response})
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
main()
|
main()
|
||||||
Loading…
Reference in a new issue