mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-03-26 17:51:54 +01:00
4fe99d05fd
* add timings for syncing * add more logging * more debugging * refactor multipass/db check out of VespaIndex * circular imports? * more debugging * add logs * various improvements * additional logs to narrow down issue * use global httpx pool for the main vespa flows in celery. Use in more places eventually. * cleanup debug logging, etc * remove debug logging * this should use the secondary index * mypy * missed some logging * review fixes * refactor get_default_document_index to use search settings * more missed logging * fix circular refs --------- Co-authored-by: Richard Kuo (Danswer) <rkuo@onyx.app> Co-authored-by: pablodanswer <pablo@danswer.ai>
58 lines
1.7 KiB
Python
58 lines
1.7 KiB
Python
import threading
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
|
|
class HttpxPool:
|
|
"""Class to manage a global httpx Client instance"""
|
|
|
|
_clients: dict[str, httpx.Client] = {}
|
|
_lock: threading.Lock = threading.Lock()
|
|
|
|
# Default parameters for creation
|
|
DEFAULT_KWARGS = {
|
|
"http2": True,
|
|
"limits": lambda: httpx.Limits(),
|
|
}
|
|
|
|
def __init__(self) -> None:
|
|
pass
|
|
|
|
@classmethod
|
|
def _init_client(cls, **kwargs: Any) -> httpx.Client:
|
|
"""Private helper method to create and return an httpx.Client."""
|
|
merged_kwargs = {**cls.DEFAULT_KWARGS, **kwargs}
|
|
return httpx.Client(**merged_kwargs)
|
|
|
|
@classmethod
|
|
def init_client(cls, name: str, **kwargs: Any) -> None:
|
|
"""Allow the caller to init the client with extra params."""
|
|
with cls._lock:
|
|
if name not in cls._clients:
|
|
cls._clients[name] = cls._init_client(**kwargs)
|
|
|
|
@classmethod
|
|
def close_client(cls, name: str) -> None:
|
|
"""Allow the caller to close the client."""
|
|
with cls._lock:
|
|
client = cls._clients.pop(name, None)
|
|
if client:
|
|
client.close()
|
|
|
|
@classmethod
|
|
def close_all(cls) -> None:
|
|
"""Close all registered clients."""
|
|
with cls._lock:
|
|
for client in cls._clients.values():
|
|
client.close()
|
|
cls._clients.clear()
|
|
|
|
@classmethod
|
|
def get(cls, name: str) -> httpx.Client:
|
|
"""Gets the httpx.Client. Will init to default settings if not init'd."""
|
|
with cls._lock:
|
|
if name not in cls._clients:
|
|
cls._clients[name] = cls._init_client()
|
|
return cls._clients[name]
|