chatbot.py•2.13 kB
import os
import sys
import json
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from openai import OpenAI
from typing import Optional
from pydantic import BaseModel
from server.prompt import QUERY_SYSTEM_PROMPT, SEC_SYSTEM_PROMPT, SEC_USER_PROMPT
from server.pc import Pinecone_DB
class Query(BaseModel):
query: str
company_ticker: Optional[str] = None
filing_type: Optional[str] = None
filing_date: Optional[str] = None
class Chatbot:
def __init__(self, model: str = 'gpt-5'):
self.client = OpenAI(api_key = os.environ.get('OPENAI_API_KEY'))
self.model = model
self.pc = Pinecone_DB()
def query_condenser(self, query: str) -> str:
response = self.client.responses.parse(
model = self.model,
input = [
{'role': 'system', 'content': QUERY_SYSTEM_PROMPT},
{'role': 'user', 'content': query}
],
text_format = Query
)
return response.output_parsed
def query_sec(self, query: Query) -> str:
pc_query = f'{query.query}'
if query.company_ticker:
pc_query += f' company: {query.company_ticker}'
if query.filing_type:
pc_query += f' filing_type: {query.filing_type}'
if query.filing_date:
pc_query += f' filing_date: {query.filing_date}'
results = self.pc.query(pc_query, top_k = 30, top_n = 30, rerank = True)
response = self.client.responses.create(
model = self.model,
input = [
{'role': 'system', 'content': SEC_SYSTEM_PROMPT},
{'role': 'user', 'content': SEC_USER_PROMPT.format(question = query, documents = results)}
]
)
return response.output_text
def run(self, query: str) -> str:
query = self.query_condenser(query)
response = self.query_sec(query)
return response
if __name__ == '__main__':
chatbot = Chatbot()
query = 'What is the latest 10-K for Apple Inc.?'
response = chatbot.run(query)
print(response)