__init__.py 901 B

1234567891011121314151617181920212223242526272829303132
  1. import os
  2. from typing import Optional
  3. import langchain
  4. from flask import Flask
  5. from langchain.prompts.base import DEFAULT_FORMATTER_MAPPING
  6. from pydantic import BaseModel
  7. from core.callback_handler.std_out_callback_handler import DifyStdOutCallbackHandler
  8. from core.prompt.prompt_template import OneLineFormatter
  9. class HostedOpenAICredential(BaseModel):
  10. api_key: str
  11. class HostedLLMCredentials(BaseModel):
  12. openai: Optional[HostedOpenAICredential] = None
  13. hosted_llm_credentials = HostedLLMCredentials()
  14. def init_app(app: Flask):
  15. formatter = OneLineFormatter()
  16. DEFAULT_FORMATTER_MAPPING['f-string'] = formatter.format
  17. if os.environ.get("DEBUG") and os.environ.get("DEBUG").lower() == 'true':
  18. langchain.verbose = True
  19. if app.config.get("OPENAI_API_KEY"):
  20. hosted_llm_credentials.openai = HostedOpenAICredential(api_key=app.config.get("OPENAI_API_KEY"))