Update app.py

app.py

@@ -26,8 +26,18 @@ import logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+import tiktoken
 
+def count_tokens(text):
+    encoding = tiktoken.get_encoding("cl100k_base")
+    return len(encoding.encode(text))
 
+def truncate_text(text, max_tokens):
+    encoding = tiktoken.get_encoding("cl100k_base")
+    encoded = encoding.encode(text)
+    if len(encoded) > max_tokens:
+        return encoding.decode(encoded[:max_tokens])
+    return text
 
 def setup_session():
     try:
@@ -126,6 +136,10 @@ def call_api(content, system_message, max_tokens, temperature, top_p, max_retrie
     for attempt in range(max_retries):
         try:
             messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
+            total_tokens = count_tokens(system_message) + count_tokens(content) + max_tokens
+            if total_tokens > 32000:
+                max_tokens = 32000 - count_tokens(system_message) - count_tokens(content) - 100  # 안전 마진
+            
             random_seed = random.randint(0, 1000000)
             response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, seed=random_seed)
             modified_text = response.choices[0].message.content
@@ -140,6 +154,15 @@ def call_api(content, system_message, max_tokens, temperature, top_p, max_retrie
                 return f"API 호출 실패: {str(e)}", 0, 0, 0
 
 
+
+
+
+
+
+
+
+
+
 def analyze_info(category, style, topic, references1, references2, references3):
     return f"선택한 카테고리: {category}\n선택한 포스팅 스타일: {style}\n블로그 주제: {topic}\n참고 글1: {references1}\n참고 글2: {references2}\n참고 글3: {references3}"
 
@@ -192,72 +215,66 @@ def process_all_titles(category, style, topic):
     return "\n".join(results)
 
 
-def generate_outline(category, style, topic, references1, references2, references3, title):
+
+
+
+
+def fetch_references(topic, max_tokens=3000):
+    search_url = generate_naver_search_url(topic)
+    session = setup_session()
+    if session is None:
+        return "Failed to set up session.", "", "", ""
+    results = crawl_naver_search_results(search_url, session)
+    if not results:
+        return "No results found.", "", "", ""
+    
+    selected_results = random.sample(results, 3)
+    references1_content = truncate_text(f"제목: {selected_results[0]['제목']}\n내용: {crawl_blog_content(selected_results[0]['링크'], session)}", max_tokens)
+    references2_content = truncate_text(f"제목: {selected_results[1]['제목']}\n내용: {crawl_blog_content(selected_results[1]['링크'], session)}", max_tokens)
+    references3_content = truncate_text(f"제목: {selected_results[2]['제목']}\n내용: {crawl_blog_content(selected_results[2]['링크'], session)}", max_tokens)
+    
+    return "참고글 생성 완료", references1_content, references2_content, references3_content
+
+
+
+def generate_outline(category, style, topic, references1, references2, references3, title, max_tokens=2000):
     full_content = analyze_info(category, style, topic, references1, references2, references3)
-    content = f"{full_content}\nTitle: {title}"
+    content = truncate_text(f"{full_content}\nTitle: {title}", max_tokens)
     category_prompt = get_outline_prompt(category)
     style_prompt = get_style_prompt(style)
-    max_tokens = 6000
-    temperature = 0.8
-    top_p = 0.95
     combined_prompt = f"{category_prompt}\n\n{style_prompt}"
-    modified_text, input_tokens, output_tokens, total_tokens = call_api(content, combined_prompt, max_tokens, temperature, top_p)
+    modified_text, input_tokens, output_tokens, total_tokens = call_api(content, combined_prompt, 6000, 0.8, 0.95)
     token_usage_message = f"[입력 토큰수: {input_tokens}]\n[출력 토큰수: {output_tokens}]\n[총 토큰수: {total_tokens}]"
     return modified_text, token_usage_message
 
-def generate_blog_post(category, style, topic, references1, references2, references3, title, outline):
+def generate_blog_post(category, style, topic, references1, references2, references3, title, outline, max_tokens=4000):
     full_content = analyze_info(category, style, topic, references1, references2, references3)
-    content = f"{full_content}\nTitle: {title}\nOutline: {outline}"
+    content = truncate_text(f"{full_content}\nTitle: {title}\nOutline: {outline}", max_tokens)
     category_prompt = get_blog_post_prompt(category)
     style_prompt = get_style_prompt(style)
-    max_tokens = 8000
-    temperature = 0.8
-    top_p = 0.95
     combined_prompt = f"{category_prompt}\n\n{style_prompt}"
-    modified_text, input_tokens, output_tokens, total_tokens = call_api(content, combined_prompt, max_tokens, temperature, top_p)
+    modified_text, input_tokens, output_tokens, total_tokens = call_api(content, combined_prompt, 8000, 0.8, 0.95)
     formatted_text = modified_text.replace('\n', '\n\n')
     token_usage_message = f"[입력 토큰수: {input_tokens}]\n[출력 토큰수: {output_tokens}]\n[총 토큰수: {total_tokens}]"
     return formatted_text, token_usage_message
 
-def fetch_references(topic):
-    search_url = generate_naver_search_url(topic)
-    session = setup_session()
-    if session is None:
-        return "Failed to set up session.", "", "", ""
-    results = crawl_naver_search_results(search_url, session)
-    if not results:
-        return "No results found.", "", "", ""
-    
-    selected_results = random.sample(results, 3)
-    references1_content = f"제목: {selected_results[0]['제목']}\n내용: {crawl_blog_content(selected_results[0]['링크'], session)}"
-    references2_content = f"제목: {selected_results[1]['제목']}\n내용: {crawl_blog_content(selected_results[1]['링크'], session)}"
-    references3_content = f"제목: {selected_results[2]['제목']}\n내용: {crawl_blog_content(selected_results[2]['링크'], session)}"
-    
-    return "참고글 생성 완료", references1_content, references2_content, references3_content
-
 def fetch_references_and_generate_all_steps(category, style, topic, blog_title):
-    search_url = generate_naver_search_url(topic)
-    session = setup_session()
-    if session is None:
-        return "", "", "", "", "", "", "", "", "", ""
-    
-    results = crawl_naver_search_results(search_url, session)
-    if not results:
-        return "", "", "", "", "", "", "", "", "", ""
-    
-    selected_results = random.sample(results, 3)
-    references1_content = f"제목: {selected_results[0]['제목']}\n내용: {crawl_blog_content(selected_results[0]['링크'], session)}"
-    references2_content = f"제목: {selected_results[1]['제목']}\n내용: {crawl_blog_content(selected_results[1]['링크'], session)}"
-    references3_content = f"제목: {selected_results[2]['제목']}\n내용: {crawl_blog_content(selected_results[2]['링크'], session)}"
+    _, references1_content, references2_content, references3_content = fetch_references(topic)
     
-    # 아웃라인 생성
     outline_result, outline_token_usage = generate_outline(category, style, topic, references1_content, references2_content, references3_content, blog_title)
     
-    # 블로그 글 생성
     blog_post_result, blog_post_token_usage = generate_blog_post(category, style, topic, references1_content, references2_content, references3_content, blog_title, outline_result)
     
     return references1_content, references2_content, references3_content, outline_result, outline_token_usage, blog_post_result, blog_post_token_usage
 
+
+
+
+
+
+
+
+
 def get_title_prompt(category):
     if (category == "일반"):
         return """