Update app.py

app.py

@@ -4,6 +4,8 @@ from transformers import GPT2Tokenizer
 import yfinance as yf
 import time
 import logging
+from datetime import datetime
+
 
 logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
 
@@ -62,6 +64,7 @@ def generate(prompt, history=[], temperature=0.1, max_new_tokens=10000, top_p=0.
     input_tokens = len(tokenizer.encode(prompt))
     total_tokens_used += input_tokens
     available_tokens = 32768 - total_tokens_used
+
     if available_tokens <= 0:
         yield f"Error: 입력이 최대 허용 토큰 수를 초과합니다. Total tokens used: {total_tokens_used}"
         return
@@ -69,23 +72,40 @@ def generate(prompt, history=[], temperature=0.1, max_new_tokens=10000, top_p=0.
     formatted_prompt = format_prompt(prompt, history)
     output_accumulated = ""
     try:
-        stream = client.text_generation(
-            formatted_prompt,
-            temperature=temperature,
-            max_new_tokens=min(max_new_tokens, available_tokens),
-            top_p=top_p,
-            repetition_penalty=repetition_penalty,
-            do_sample=True,
-            seed=42,
-            stream=True
-        )
+        stream = client.text_generation(formatted_prompt, temperature=temperature, max_new_tokens=min(max_new_tokens, available_tokens),
+                                        top_p=top_p, repetition_penalty=repetition_penalty, do_sample=True, seed=42, stream=True)
         for response in stream:
             output_part = response['generated_text'] if 'generated_text' in response else str(response)
             output_accumulated += output_part
-            yield process_financial_data(output_part) + f"\n\n---\nTotal tokens used: {total_tokens_used}"
+            
+            # 데이터 다운로드 및 결과 출력
+            if "티커" in output_part:
+                ticker = extract_ticker(output_part)  # 티커 추출하는 더미 함수, 실제로는 티커를 추출하는 로직 필요
+                download_result = download_stock_data(ticker)
+                output_accumulated += download_result
+            
+            yield output_accumulated + f"\n\n---\nTotal tokens used: {total_tokens_used}"
     except Exception as e:
         yield f"Error: {str(e)}\nTotal tokens used: {total_tokens_used}"
 
+def download_stock_data(ticker):
+    try:
+        today = datetime.now()
+        start_date = today.replace(year=today.year - 10).strftime('%Y-%m-%d')
+        end_date = today.strftime('%Y-%m-%d')
+        data = yf.download(ticker, start=start_date, end=end_date)
+        if data.empty:
+            return f"Error: 데이터를 다운로드할 수 없습니다. 티커 {ticker}를 확인하세요."
+        else:
+            return f"Success: {ticker} 데이터 다운로드 성공."
+    except Exception as e:
+        return f"Error: 데이터 다운로드 중 오류가 발생했습니다. {str(e)}"
+
+def extract_ticker(output_part):
+    # 이 함수는 사용자 출력에서 티커를 추출하는 로직을 구현해야 함
+    # 여기서는 예시로 'AAPL'을 반환하도록 설정
+    return "AAPL"
+
 def validate_ticker(ticker):
     logging.debug(f"Validating ticker: {ticker}")
     stock = yf.Ticker(ticker)