testr2r.py 2.8 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485
  1. from typing import Optional, Any
  2. from r2r import R2RClient
  3. from app.libs.util import verify_jwt_expiration
  4. from config.llm import tool_settings
  5. import inspect
  6. import r2r
  7. print(r2r.__version__)
  8. class R2R:
  9. client: R2RClient
  10. def init(self):
  11. self.client = R2RClient(tool_settings.R2R_BASE_URL)
  12. self.auth_enabled = tool_settings.R2R_USERNAME and tool_settings.R2R_PASSWORD
  13. self._login()
  14. def ingest_file(self, file_path: str, metadata: Optional[dict]):
  15. self._check_login()
  16. ingest_response = self.client.ingest_files(
  17. file_paths=[file_path],
  18. metadatas=[metadata] if metadata else None
  19. )
  20. return ingest_response.get("results")
  21. def search(self, query: str, filters: dict[str, Any]):
  22. self._check_login()
  23. search_response = self.client.search(
  24. query=query,
  25. vector_search_settings={
  26. "filters": filters,
  27. "search_limit": tool_settings.R2R_SEARCH_LIMIT,
  28. "do_hybrid_search": True
  29. },
  30. )
  31. return search_response.get("results").get("vector_search_results")
  32. def _login(self):
  33. if not self.auth_enabled:
  34. return
  35. print("test")
  36. print(self.client)
  37. print(self.client.users)
  38. print(self.client._wrap_v3_methods)
  39. # 获取实例方法
  40. instance_methods = inspect.getmembers(self.client._wrap_v3_methods, predicate=inspect.isfunction)
  41. print("实例方法:", [method[0] for method in instance_methods])
  42. # 获取类方法(包括静态方法)
  43. class_methods = inspect.getmembers(self.client._wrap_v3_methods, predicate=inspect.ismethod)
  44. print("类方法:", [method[0] for method in class_methods])
  45. #self.client.user.login(tool_settings.R2R_USERNAME, tool_settings.R2R_PASSWORD)
  46. def _check_login(self):
  47. if not self.auth_enabled:
  48. return
  49. if verify_jwt_expiration(self.client.access_token):
  50. return
  51. if verify_jwt_expiration(self.client._refresh_token):
  52. self.client.refresh_access_token()
  53. else:
  54. self._login()
  55. def get_method_parameters(method):
  56. params = inspect.signature(method).parameters
  57. for name, param in params.items():
  58. default_value = param.default if param.default is not param.empty else None
  59. print(f"参数名: {name}, 默认值: {default_value}")
  60. # 获取实例方法
  61. instance_methods = inspect.getmembers(R2RClient, predicate=inspect.isfunction)
  62. print("实例方法:", [method[0] for method in instance_methods])
  63. # 获取类方法(包括静态方法)
  64. class_methods = inspect.getmembers(R2RClient, predicate=inspect.ismethod)
  65. print("类方法:", [method[0] for method in class_methods])
  66. r2r1 = R2R()
  67. r2r1.init()