src/get_stock_code.py

import os
from openai import AzureOpenAI
from dotenv import load_dotenv
from langchain.document_loaders import UnstructuredCSVLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma
from langchain.chat_models import ChatOpenAI
from langchain.chains import ConversationalRetrievalChain
from langchain.text_splitter import RecursiveCharacterTextSplitter
import pandas as pd


FILE_PATH = "./data/csv/edinet_df.csv"

load_dotenv()
openai_endpoint = os.environ.get("AZURE_OPENAI_ENDPOINT")
openai_api_key = os.environ.get("AZURE_OPENAI_API_KEY")
version = os.environ.get("AZURE_OPENAI_API_VERSION")
api_type = 'gpt-4o'

class StockCode_getter:

    def __init__(self):
        self.client = AzureOpenAI(
            api_key=openai_api_key,
            api_version=version,
            azure_endpoint=openai_endpoint
        )
        df = pd.read_csv(FILE_PATH)
        # 特定の列を抽出する（例として 'column1' と 'column3' を抽出）
        selected_columns = df[['会社名', '証券コード']]

        # 会社名と証券コードのみのcsvファイルを作成
        selected_columns.to_csv('./data/csv/selected_columns.csv', index=False)
        self.company_file = './data/csv/selected_columns.csv'
        # # load Excel file
        # loader = UnstructuredCSVLoader(self.company_file)
        # docs = loader.load()
        # # ドキュメントを小さなチャンクに分割
        # text_splitter = RecursiveCharacterTextSplitter(chunk_size=30, chunk_overlap=10)
        # docs = text_splitter.split_documents(docs)

        # # create vector store
        # self.vector_stores = Chroma.from_documents(docs, OpenAIEmbeddings())

    def get_stock_code(self, company_name):
        df = pd.read_csv('./data/csv/edinet_df.csv')
        filtered_df = df[df['会社名'].str.contains(company_name)]
        if not filtered_df.empty:
            company = filtered_df.iloc[0]
            messages = [
            {"role": "system", "content":
                f"""You are a knowledgeable assistant with expertise in corporate information and financial markets.
                When users inquire about the listing status of a company, provide a clear Yes or No answer."""},
            {"role": "user", "content":f"""{company_name}は以下の企業リストに含まれますか？ Yes/No で答えてください。
                以下のドキュメントには企業名と証券コードが記載されています。企業名があればYesと答えてください。なければNoと答えてください。
                ===ドキュメント===
                {company}

                ===Example1===
                [ドキュメント]
                会社名  株式会社SONY\n証券コード 67580
                [input]
                株式会社SONY
                [output]
                Yes

                ===Example2===
                [ドキュメント]
                会社名  株式会社SONY\n証券コード 67580
                [input]
                株式会社Almondo
                [output]
                No

                ===Let's begin===
                [input]
                {company_name}
                [output]

                """}
            ]

            response = self.client.chat.completions.create(
                    messages=messages,
                    model=api_type,
                    temperature=0,
            )
            is_listed = (response.choices[0].message.content)

            if is_listed == "Yes":
                docID = company['docID']
                return docID
        return -1