hosted.py 2.7 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576
  1. import os
  2. from typing import Optional
  3. import langchain
  4. from flask import Flask
  5. from pydantic import BaseModel
  6. class HostedOpenAI(BaseModel):
  7. api_base: str = None
  8. api_organization: str = None
  9. api_key: str
  10. quota_limit: int = 0
  11. """Quota limit for the openai hosted model. 0 means unlimited."""
  12. paid_enabled: bool = False
  13. paid_stripe_price_id: str = None
  14. paid_increase_quota: int = 1
  15. class HostedAzureOpenAI(BaseModel):
  16. api_base: str
  17. api_key: str
  18. quota_limit: int = 0
  19. """Quota limit for the azure openai hosted model. 0 means unlimited."""
  20. class HostedAnthropic(BaseModel):
  21. api_base: str = None
  22. api_key: str
  23. quota_limit: int = 0
  24. """Quota limit for the anthropic hosted model. 0 means unlimited."""
  25. paid_enabled: bool = False
  26. paid_stripe_price_id: str = None
  27. paid_increase_quota: int = 1
  28. class HostedModelProviders(BaseModel):
  29. openai: Optional[HostedOpenAI] = None
  30. azure_openai: Optional[HostedAzureOpenAI] = None
  31. anthropic: Optional[HostedAnthropic] = None
  32. hosted_model_providers = HostedModelProviders()
  33. def init_app(app: Flask):
  34. if os.environ.get("DEBUG") and os.environ.get("DEBUG").lower() == 'true':
  35. langchain.verbose = True
  36. if app.config.get("HOSTED_OPENAI_ENABLED"):
  37. hosted_model_providers.openai = HostedOpenAI(
  38. api_base=app.config.get("HOSTED_OPENAI_API_BASE"),
  39. api_organization=app.config.get("HOSTED_OPENAI_API_ORGANIZATION"),
  40. api_key=app.config.get("HOSTED_OPENAI_API_KEY"),
  41. quota_limit=app.config.get("HOSTED_OPENAI_QUOTA_LIMIT"),
  42. paid_enabled=app.config.get("HOSTED_OPENAI_PAID_ENABLED"),
  43. paid_stripe_price_id=app.config.get("HOSTED_OPENAI_PAID_STRIPE_PRICE_ID"),
  44. paid_increase_quota=app.config.get("HOSTED_OPENAI_PAID_INCREASE_QUOTA"),
  45. )
  46. if app.config.get("HOSTED_AZURE_OPENAI_ENABLED"):
  47. hosted_model_providers.azure_openai = HostedAzureOpenAI(
  48. api_base=app.config.get("HOSTED_AZURE_OPENAI_API_BASE"),
  49. api_key=app.config.get("HOSTED_AZURE_OPENAI_API_KEY"),
  50. quota_limit=app.config.get("HOSTED_AZURE_OPENAI_QUOTA_LIMIT"),
  51. )
  52. if app.config.get("HOSTED_ANTHROPIC_ENABLED"):
  53. hosted_model_providers.anthropic = HostedAnthropic(
  54. api_base=app.config.get("HOSTED_ANTHROPIC_API_BASE"),
  55. api_key=app.config.get("HOSTED_ANTHROPIC_API_KEY"),
  56. quota_limit=app.config.get("HOSTED_ANTHROPIC_QUOTA_LIMIT"),
  57. paid_enabled=app.config.get("HOSTED_ANTHROPIC_PAID_ENABLED"),
  58. paid_stripe_price_id=app.config.get("HOSTED_ANTHROPIC_PAID_STRIPE_PRICE_ID"),
  59. paid_increase_quota=app.config.get("HOSTED_ANTHROPIC_PAID_INCREASE_QUOTA"),
  60. )