@@ -27,13 +27,25 @@ def __new__(cls, value):
2727
2828
2929class LangTraceFile (io .BytesIO ):
30- _host : str = os .environ .get ("LANGTRACE_API_HOST" , None ) or LANGTRACE_REMOTE_URL
3130
3231 def __init__ (self , fs : "LangTraceFileSystem" , path : str , mode : OpenMode ):
3332 super ().__init__ ()
3433 self .fs = fs
3534 self .path = path
3635 self .mode = mode
36+ self ._host : str = os .environ .get ("LANGTRACE_API_HOST" , LANGTRACE_REMOTE_URL )
37+ self ._api_key : str = os .environ .get ("LANGTRACE_API_KEY" , None )
38+ if self ._host .endswith ("/api/trace" ):
39+ self ._host = self ._host .replace ("/api/trace" , "" )
40+
41+ if self ._api_key is None :
42+ print (Fore .RED )
43+ print (
44+ f"Missing Langtrace API key, proceed to { self ._host } to create one"
45+ )
46+ print ("Set the API key as an environment variable LANGTRACE_API_KEY" )
47+ print (Fore .RESET )
48+ return
3749
3850 def close (self ) -> None :
3951 if not self .closed :
@@ -71,7 +83,7 @@ def upload_to_server(self, file_data: bytes) -> None:
7183 data = json .dumps (data ),
7284 headers = {
7385 "Content-Type" : "application/json" ,
74- "x-api-key" : os . environ . get ( "LANGTRACE_API_KEY" ) ,
86+ "x-api-key" : self . _api_key ,
7587 },
7688 timeout = 20 ,
7789 )
@@ -82,14 +94,26 @@ def upload_to_server(self, file_data: bytes) -> None:
8294
8395
8496class LangTraceFileSystem (AbstractFileSystem ):
85- _host : str = os .environ .get ("LANGTRACE_API_HOST" , None ) or LANGTRACE_REMOTE_URL
8697 protocol = "langtracefs"
8798 sep = "/"
8899
89100 def __init__ (self , * args , ** kwargs ):
90101 super ().__init__ (* args , ** kwargs )
91102 self .files = {}
92103 self .dirs = set ()
104+ self ._host : str = os .environ .get ("LANGTRACE_API_HOST" , LANGTRACE_REMOTE_URL )
105+ self ._api_key : str = os .environ .get ("LANGTRACE_API_KEY" , None )
106+ if self ._host .endswith ("/api/trace" ):
107+ self ._host = self ._host .replace ("/api/trace" , "" )
108+
109+ if self ._api_key is None :
110+ print (Fore .RED )
111+ print (
112+ f"Missing Langtrace API key, proceed to { self ._host } to create one"
113+ )
114+ print ("Set the API key as an environment variable LANGTRACE_API_KEY" )
115+ print (Fore .RESET )
116+ return
93117
94118 def open (
95119 self ,
@@ -118,7 +142,7 @@ def fetch_file_from_api(self, dataset_id: str) -> bytes:
118142 url = f"{ self ._host } /api/dataset/download?id={ dataset_id } " ,
119143 headers = {
120144 "Content-Type" : "application/json" ,
121- "x-api-key" : os . environ . get ( "LANGTRACE_API_KEY" ) ,
145+ "x-api-key" : self . _api_key ,
122146 },
123147 timeout = 20 ,
124148 )
0 commit comments