response.py 3.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  1. from pydantic import BaseModel
  2. from typing import List,Optional,Any,Dict
  3. from .request import TextData
  4. class Response(BaseModel):
  5. status: str
  6. ################################################
  7. class VectorData(BaseModel):
  8. text: str
  9. vector: List[float]
  10. class EmbeddingResponse(Response):
  11. vectors: List[VectorData]
  12. ################################################
  13. class SearchResult(BaseModel):
  14. id: str
  15. title: str
  16. author: str
  17. site_name: str
  18. site_domain: str
  19. pub_date: str
  20. article_text: str
  21. score: float
  22. url: str
  23. class SearchResultMeta(BaseModel):
  24. id: str
  25. title: str
  26. pub_date: str
  27. url: str
  28. class SearchResultContent(BaseModel):
  29. id: str
  30. title: str
  31. article_text: str
  32. class SearchResponse(Response):
  33. results: List[SearchResult]
  34. class SearchResponseV2(BaseModel):
  35. meta: List[SearchResultMeta]
  36. content: List[SearchResultContent]
  37. ################################################
  38. class GetDocumentResponse(Response):
  39. title: str
  40. text: str
  41. url: Optional[str]
  42. class GetMDocumentResponse(Response):
  43. results: List[GetDocumentResponse]
  44. ################################################
  45. ################################################
  46. SUCCESS = 200
  47. FAILED = 500
  48. PROCESSING = 202
  49. class StandardResponse(BaseModel):
  50. code: int = 200
  51. message: str = "success"
  52. meta: Dict = {}
  53. records: List[Any] = []
  54. def __init__(self, code=SUCCESS, success=True, err_code=-1, err_message="", meta={}, records=None, message="success"):
  55. super().__init__()
  56. self.code = code
  57. self.message = message
  58. self.meta = meta
  59. if records:
  60. recordsItem = []
  61. for item in records:
  62. #遍历item的属性,如果属性是list,就把list中的每个元素都转换成dict
  63. if isinstance(item, dict):
  64. recordsItem.append(item)
  65. continue
  66. item_dict = {}
  67. for key in item.__dict__:
  68. if not key.startswith("_"):
  69. value = item.__dict__[key]
  70. if isinstance(value, str):
  71. item_dict[key] = value
  72. elif isinstance(value, list):
  73. values = []
  74. for v in value:
  75. values.append(v)
  76. item_dict[key] = values
  77. else:
  78. item_dict[key] = value
  79. recordsItem.append(item_dict)
  80. self.records = recordsItem
  81. if success == False:
  82. self.code = FAILED
  83. if err_code >=0:
  84. self.code = err_code
  85. self.message = err_message
  86. class TaskResponse(StandardResponse):
  87. task_id: str
  88. status: str
  89. result: Optional[Any] = None