app.py

import discord
import logging
import os
import asyncio
import subprocess
from huggingface_hub import InferenceClient
from yahooquery import Ticker
from datetime import datetime

# 로깅 설정
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])

# 인텐트 설정
intents = discord.Intents.default()
intents.message_content = True
intents.messages = True
intents.guilds = True
intents.guild_messages = True

# 추론 API 클라이언트 설정
hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))

# 특정 채널 ID
SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))

# 대화 히스토리를 저장할 전역 변수
conversation_history = []

class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False

    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')
        subprocess.Popen(["python", "web.py"])
        logging.info("web.py 서버가 시작되었습니다.")

    async def on_message(self, message):
        if message.author == self.user:
            return
        if not self.is_message_in_specific_channel(message):
            return
        if self.is_processing:
            return
        self.is_processing = True
        try:
            response = await handle_message(message)
            await message.channel.send(response)
        finally:
            self.is_processing = False

    def is_message_in_specific_channel(self, message):
        return message.channel.id == SPECIFIC_CHANNEL_ID or (
            isinstance(message.channel, discord.Thread) and message.channel.parent_id == SPECIFIC_CHANNEL_ID
        )

async def handle_message(message):
    user_input = message.content
    user_mention = message.author.mention

    # 티커를 제공하지 않은 경우
    if not user_input.startswith("!t "):
        ticker = await find_ticker(user_input)
        if ticker:
            return f"{user_mention}, 종목명을 찾았습니다: {ticker}. 티커 정보를 원하면 '!t {ticker}'라고 입력하세요."
        else:
            return f"{user_mention}, 종목명을 찾을 수 없습니다. 올바른 종목명을 입력하세요."

    # 티커를 제공한 경우
    ticker = user_input[3:].strip().upper()
    return await handle_ticker(ticker, user_mention)

async def handle_ticker(ticker, user_mention):
    global conversation_history  # 전역 변수 사용을 명시
    conversation_history.append({"role": "user", "content": ticker})
    logging.debug(f'Conversation history updated: {conversation_history}')

    company_info, charts, news_summary, prediction = await get_stock_info(ticker)

    conversation_history.append({"role": "assistant", "content": prediction})
    if len(conversation_history) > 20:
        conversation_history.pop(0)

    return f"""
    **회사 정보:**
    회사 이름: {company_info.get('longName', 'N/A')}
    티커: {ticker}
    산업: {company_info.get('industry', 'N/A')}
    섹터: {company_info.get('sector', 'N/A')}
    시가 총액: {company_info.get('marketCap', 'N/A')}
    현재가: {company_info.get('currentPrice', 'N/A')}
    데이터 업데이트: {datetime.now().strftime("%Y-%m-%d %H:%M:%S")}

    **최근 1일 차트:**
    {charts['1d']}
    **최근 1주 차트:**
    {charts['1w']}
    **최근 1개월 차트:**
    {charts['1m']}
    **최근 1년 차트:**
    {charts['1y']}
    **최근 10년 차트:**
    {charts['10y']}

    **최신 뉴스:**
    {news_summary}

    **향후 1주일간 시세 예측:** 
    {prediction}
    """

async def find_ticker(query):
    tickers = Ticker(query)
    for ticker in tickers.symbols:
        return ticker
    return None

async def get_stock_info(ticker):
    tickers = Ticker(ticker)
    summary = tickers.summary_detail[ticker]
    price = summary.get("previousClose", "N/A")

    company_info = {
        'longName': tickers.asset_profile[ticker].get('longName', 'N/A'),
        'industry': tickers.asset_profile[ticker].get('industry', 'N/A'),
        'sector': tickers.asset_profile[ticker].get('sector', 'N/A'),
        'marketCap': summary.get('marketCap', 'N/A'),
        'currentPrice': price,
    }

    charts = {
        '1d': get_chart_summary(tickers.history(period="1d", interval="1d")),
        '1w': get_chart_summary(tickers.history(period="5d", interval="1d")),
        '1m': get_chart_summary(tickers.history(period="1mo", interval="1d")),
        '1y': get_chart_summary(tickers.history(period="1y", interval="1d")),
        '10y': get_chart_summary(tickers.history(period="10y", interval="1d")),
    }

    news_summary = get_news_summary(tickers.news())

    prediction = await predict_stock(ticker)

    return company_info, charts, news_summary, prediction

def get_chart_summary(history):
    if history.empty:
        return "데이터 없음"
    avg_price = history['close'].mean()
    max_price = history['close'].max()
    min_price = history['close'].min()
    return f"평균가: {avg_price:.2f}, 최고가: {max_price:.2f}, 최저가: {min_price:.2f}"

def get_news_summary(news):
    if not news:
        return "뉴스 요약을 가져오는 데 실패했습니다."
    news_summary = ""
    for article in news[:5]:
        news_summary += f"- {article['title']}: {article['summary']}\n"
    return news_summary.strip()

async def predict_stock(ticker):
    messages = [{"role": "system", "content": f"사용자가 티커 '{ticker}'에 대해 알고 싶어합니다."}] + conversation_history
    logging.debug(f'모델에 보낼 메시지: {messages}')

    loop = asyncio.get_event_loop()
    response = await loop.run_in_executor(None, lambda: hf_client.chat_completion(
        messages, max_tokens=1000, stream=True, temperature=0.7, top_p=0.85))

    full_response = []
    for part in response:
        if part.choices and part.choices[0].delta and part.choices[0].delta.content:
            full_response.append(part.choices[0].delta.content)

    return ''.join(full_response).strip()

if __name__ == "__main__":
    discord_token = os.getenv('DISCORD_TOKEN')
    discord_client = MyClient(intents=intents)
    discord_client.run(discord_token)