Update app.py

app.py

@@ -1,19 +1,82 @@
-from flask import Flask, request
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from flask import Flask, request, jsonify
+from typing import Dict, Any
+import openai
+
+openai.api_key = "YOUR_API_KEY"
+
+def submit_message(user_token, prompt, prompt_template, temperature, max_tokens, state):
+    history = state['messages']
+
+    if not prompt:
+        return {
+            'chat_messages': [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)],
+            'total_tokens_used_msg': f"Total tokens used: {state['total_tokens']} / 3000"
+        }
+
+    prompt_template = prompt_templates[prompt_template]
+
+    system_prompt = []
+    if prompt_template:
+        system_prompt = [{ "role": "system", "content": prompt_template }]
+
+    prompt_msg = { "role": "user", "content": prompt }
+    
+    try:
+        response = openai.Completion.create(
+            engine="davinci",
+            prompt=system_prompt + history + [prompt_msg],
+            temperature=temperature,
+            max_tokens=max_tokens
+        )
+
+        history.append(prompt_msg)
+        history.append({
+            "role": "system",
+            "content": response.choices[0].text
+        })
+
+        state['total_tokens'] += response['total_characters']
+    
+    except Exception as e:
+        history.append(prompt_msg)
+        history.append({
+            "role": "system",
+            "content": f"Error: {e}"
+        })
+
+    total_tokens_used_msg = f"Total tokens used: {state['total_tokens']} / 3000" if not user_token else ""
+    chat_messages = [(history[i]['content'], history[i+1]['content']) for i in range(0, len(history)-1, 2)]
+    input_visibility
+
+    return {
+        'chat_messages': chat_messages,
+        'total_tokens_used_msg': total_tokens_used_msg
+    }
 
 app = Flask(__name__)
 
-model_name = "openai-gpt"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-@app.route('/generate_text', methods=['POST'])
-def generate_text():
-    input_text = request.json['input_text']
-    input_ids = tokenizer.encode(input_text, return_tensors='pt')
-    output_ids = model(**input_ids)
-    output_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    return {'generated_text': output_text}
+state = {
+    'messages': [],
+    'total_tokens': 0
+}
+
+@app.route('/chat', methods=['POST'])
+def chat():
+    data = request.json
+    user_token = data.get('user_token')
+    prompt = data.get('prompt')
+    prompt_template = data.get('prompt_template')
+    temperature = data.get('temperature', 0.5)
+    max_tokens = data.get('max_tokens', 50)
+
+    _, chat_messages, total_tokens_used_msg, state = submit_message(user_token, prompt, prompt_template, temperature, max_tokens, state)
+
+    response = {
+        'chat_messages': chat_messages,
+        'total_tokens_used_msg': total_tokens_used_msg
+    }
+
+    return jsonify(response)
 
 if __name__ == '__main__':
     app.run()