forked from Archives/langchain
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
224 lines
6.9 KiB
Python
224 lines
6.9 KiB
Python
import argparse
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
from typing import Generator, List, Optional
|
|
|
|
import requests
|
|
import yaml
|
|
|
|
from langchain.env import get_runtime_environment
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_DIR = Path(__file__).parent
|
|
|
|
|
|
def get_docker_compose_command() -> List[str]:
|
|
if shutil.which("docker-compose") is None:
|
|
return ["docker", "compose"]
|
|
else:
|
|
return ["docker-compose"]
|
|
|
|
|
|
def get_ngrok_url(auth_token: Optional[str]) -> str:
|
|
"""Get the ngrok URL for the LangChainPlus server."""
|
|
ngrok_url = "http://localhost:4040/api/tunnels"
|
|
try:
|
|
response = requests.get(ngrok_url)
|
|
response.raise_for_status()
|
|
exposed_url = response.json()["tunnels"][0]["public_url"]
|
|
except requests.exceptions.HTTPError:
|
|
raise ValueError("Could not connect to ngrok console.")
|
|
except (KeyError, IndexError):
|
|
message = "ngrok failed to start correctly. "
|
|
if auth_token is not None:
|
|
message += "Please check that your authtoken is correct."
|
|
raise ValueError(message)
|
|
return exposed_url
|
|
|
|
|
|
@contextmanager
|
|
def create_ngrok_config(
|
|
auth_token: Optional[str] = None,
|
|
) -> Generator[Path, None, None]:
|
|
"""Create the ngrok configuration file."""
|
|
config_path = _DIR / "ngrok_config.yaml"
|
|
if config_path.exists():
|
|
# If there was an error in a prior run, it's possible
|
|
# Docker made this a directory instead of a file
|
|
if config_path.is_dir():
|
|
shutil.rmtree(config_path)
|
|
else:
|
|
config_path.unlink()
|
|
ngrok_config = {
|
|
"tunnels": {
|
|
"langchain": {
|
|
"proto": "http",
|
|
"addr": "langchain-backend:8000",
|
|
}
|
|
},
|
|
"version": "2",
|
|
"region": "us",
|
|
}
|
|
if auth_token is not None:
|
|
ngrok_config["authtoken"] = auth_token
|
|
config_path = _DIR / "ngrok_config.yaml"
|
|
with config_path.open("w") as f:
|
|
yaml.dump(ngrok_config, f)
|
|
yield config_path
|
|
# Delete the config file after use
|
|
config_path.unlink(missing_ok=True)
|
|
|
|
|
|
class ServerCommand:
|
|
"""Manage the LangChainPlus Tracing server."""
|
|
|
|
def __init__(self) -> None:
|
|
self.docker_compose_command = get_docker_compose_command()
|
|
self.docker_compose_file = (
|
|
Path(__file__).absolute().parent / "docker-compose.yaml"
|
|
)
|
|
self.ngrok_path = Path(__file__).absolute().parent / "docker-compose.ngrok.yaml"
|
|
|
|
def _start_local(self) -> None:
|
|
command = [
|
|
*self.docker_compose_command,
|
|
"-f",
|
|
str(self.docker_compose_file),
|
|
]
|
|
subprocess.run(
|
|
[
|
|
*command,
|
|
"up",
|
|
"--pull=always",
|
|
"--quiet-pull",
|
|
"--wait",
|
|
]
|
|
)
|
|
logger.info(
|
|
"LangChain server is running at http://localhost. To connect"
|
|
" locally, set the following environment variable"
|
|
" when running your LangChain application."
|
|
)
|
|
|
|
logger.info("\tLANGCHAIN_TRACING_V2=true")
|
|
subprocess.run(["open", "http://localhost"])
|
|
|
|
def _start_and_expose(self, auth_token: Optional[str]) -> None:
|
|
with create_ngrok_config(auth_token=auth_token):
|
|
command = [
|
|
*self.docker_compose_command,
|
|
"-f",
|
|
str(self.docker_compose_file),
|
|
"-f",
|
|
str(self.ngrok_path),
|
|
]
|
|
subprocess.run(
|
|
[
|
|
*command,
|
|
"up",
|
|
"--pull=always",
|
|
"--quiet-pull",
|
|
"--wait",
|
|
]
|
|
)
|
|
logger.info(
|
|
"ngrok is running. You can view the dashboard at http://0.0.0.0:4040"
|
|
)
|
|
ngrok_url = get_ngrok_url(auth_token)
|
|
logger.info(
|
|
"LangChain server is running at http://localhost."
|
|
" To connect remotely, set the following environment"
|
|
" variable when running your LangChain application."
|
|
)
|
|
logger.info("\tLANGCHAIN_TRACING_V2=true")
|
|
logger.info(f"\tLANGCHAIN_ENDPOINT={ngrok_url}")
|
|
subprocess.run(["open", "http://localhost"])
|
|
|
|
def start(self, *, expose: bool = False, auth_token: Optional[str] = None) -> None:
|
|
"""Run the LangChainPlus server locally.
|
|
|
|
Args:
|
|
expose: If True, expose the server to the internet using ngrok.
|
|
auth_token: The ngrok authtoken to use (visible in the ngrok dashboard).
|
|
If not provided, ngrok server session length will be restricted.
|
|
"""
|
|
|
|
if expose:
|
|
self._start_and_expose(auth_token=auth_token)
|
|
else:
|
|
self._start_local()
|
|
|
|
def stop(self) -> None:
|
|
"""Stop the LangChainPlus server."""
|
|
subprocess.run(
|
|
[
|
|
*self.docker_compose_command,
|
|
"-f",
|
|
str(self.docker_compose_file),
|
|
"-f",
|
|
str(self.ngrok_path),
|
|
"down",
|
|
]
|
|
)
|
|
|
|
|
|
def env() -> None:
|
|
"""Print the runtime environment information."""
|
|
env = get_runtime_environment()
|
|
logger.info("LangChain Environment:")
|
|
logger.info("\n".join(f"{k}:{v}" for k, v in env.items()))
|
|
|
|
|
|
def main() -> None:
|
|
"""Main entrypoint for the CLI."""
|
|
parser = argparse.ArgumentParser()
|
|
subparsers = parser.add_subparsers(description="LangChainPlus CLI commands")
|
|
|
|
server_command = ServerCommand()
|
|
server_parser = subparsers.add_parser("server", description=server_command.__doc__)
|
|
server_subparsers = server_parser.add_subparsers()
|
|
|
|
server_start_parser = server_subparsers.add_parser(
|
|
"start", description="Start the LangChainPlus server."
|
|
)
|
|
server_start_parser.add_argument(
|
|
"--expose",
|
|
action="store_true",
|
|
help="Expose the server to the internet using ngrok.",
|
|
)
|
|
server_start_parser.add_argument(
|
|
"--ngrok-authtoken",
|
|
default=os.getenv("NGROK_AUTHTOKEN"),
|
|
help="The ngrok authtoken to use (visible in the ngrok dashboard)."
|
|
" If not provided, ngrok server session length will be restricted.",
|
|
)
|
|
server_start_parser.set_defaults(
|
|
func=lambda args: server_command.start(
|
|
expose=args.expose, auth_token=args.ngrok_authtoken
|
|
)
|
|
)
|
|
|
|
server_stop_parser = server_subparsers.add_parser(
|
|
"stop", description="Stop the LangChainPlus server."
|
|
)
|
|
server_stop_parser.set_defaults(func=lambda args: server_command.stop())
|
|
|
|
env_parser = subparsers.add_parser("env")
|
|
env_parser.set_defaults(func=lambda args: env())
|
|
|
|
args = parser.parse_args()
|
|
if not hasattr(args, "func"):
|
|
parser.print_help()
|
|
return
|
|
args.func(args)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|