File size: 4,774 Bytes
03b26f3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
91b0ac2
03b26f3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
91b0ac2
03b26f3
 
 
 
91b0ac2
03b26f3
 
91b0ac2
 
03b26f3
 
91b0ac2
03b26f3
 
91b0ac2
 
 
 
 
 
03b26f3
91b0ac2
03b26f3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
91b0ac2
03b26f3
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
import os
from openai import OpenAI
import streamlit as st

api_key = os.getenv("NVIDIA_API_KEY")

# Check if the API key is found
if api_key is None:
    st.error("NVIDIA_API_KEY environment variable not found.")
else:
    # Initialize the OpenAI client
    client = OpenAI(
        base_url="https://integrate.api.nvidia.com/v1",
        api_key=api_key
    )

class ConversationManager:
    def __init__(self):
        # (No need to initialize history here, it will be handled in main())
        pass

    def generate_ai_response(self, prompt):
        """Generates a response from an AI model

        Args:
        prompt: The prompt to send to the AI model.

        Returns:
        response from the AI model.
        """
        try:
            # Access conversation_history from session state
            messages = [
                {
                    "role": "system",
                    "content": "You are a programming assistant focused on providing \
                    accurate, clear, and concise answers to technical questions. \
                    Your goal is to help users solve programming problems efficiently, \
                    explain concepts clearly, and provide examples when appropriate. \
                    Use a professional yet approachable tone. Use explicit markdown \
                    format for code for all codes in the output."
                }
            ]

            for message in st.session_state.conversation_manager.conversation_history:
                messages.append(message)

            messages.append({
                "role": "user",
                "content": prompt
            })

            completion = client.chat.completions.create(
                model="meta/llama-3.1-405b-instruct",
                temperature=0.5,  # Adjust temperature for creativity
                top_p=1,
                max_tokens=1024,
                messages=messages,
                stream=False
            )

            model_response = completion.choices[0].message.content

            st.session_state.conversation_manager.conversation_history.append({
                "role": "assistant",
                "content": model_response
            })

            st.session_state.conversation_manager.conversation_history.append({
                "role": "assistant",
                "content": completion.choices[0].message.content
            })                
            return model_response
            
        except Exception as e:
            st.error(f"Error handling AI response: {e}")
            return None


def main():
    # Initialize ConversationManager in session state if not already present
    if "conversation_manager" not in st.session_state:
        st.session_state.conversation_manager = ConversationManager()
        st.session_state.conversation_manager.conversation_history = []

    st.title("AI-Assisted Code Generator")

    tab1, tab2, tab3 = st.tabs(["About", "Code Generation",  "Conversation History"])

    with tab1:
        st.header("About this App")
        st.write("This app demonstrates how to use AI to assist in code generation.")

    with tab2:
        st.header("Generate Code")
        framework = st.selectbox("Select a framework", ["Streamlit", "Gradio"])
        app_details = st.text_area("Describe the app you want to create", value="Create a complete app that ")

        if st.button("Generate Prompt"):
            user_prompt = f"Using {framework}, {app_details}"
            st.write("**Generated Prompt:**", user_prompt)

            with st.spinner("Thinking..."):
                # Add the user message to the history FIRST
                st.session_state.conversation_manager.conversation_history.append({"role": "user", "content": user_prompt})  

                ai_response = st.session_state.conversation_manager.generate_ai_response(user_prompt)
                if ai_response:
                    st.session_state.conversation_manager.conversation_history.append({"role": "assistant", "content": ai_response})
                    st.markdown(f"**User:** {user_prompt}")  
                    st.markdown(f"**AI:** {ai_response}")
                else:
                    st.write("**Error:** Failed to generate AI response.")

    with tab3:
        st.header("Conversation History")
        if st.session_state.conversation_manager.conversation_history:
            for msg in st.session_state.conversation_manager.conversation_history:
                if msg['role'] == 'user':
                    st.markdown(f"**User:** {msg['content']}")
                else:
                    st.markdown(f"**AI:** {msg['content']}")
        else:
            st.write("No conversation history yet.")

if __name__ == "__main__":
    main()