diff --git a/api/main.py b/api/main.py index 92b6305..f821b70 100644 --- a/api/main.py +++ b/api/main.py @@ -1,36 +1,29 @@ -from typing import Dict, List, TypedDict import re -import uvicorn +from typing import Dict, List, TypedDict import torch +import uvicorn from fastapi import FastAPI from fastapi.staticfiles import StaticFiles - from pydantic import BaseModel -from env import settings - -from core.prompts.error import ERROR_PROMPT from core.agents.manager import AgentManager +from core.handlers.base import BaseHandler, FileHandler, FileType +from core.handlers.dataframe import CsvToDataframe +from core.handlers.image import ImageCaptioning +from core.prompts.error import ERROR_PROMPT from core.tools.base import BaseToolSet -from core.tools.terminal import Terminal +from core.tools.cpu import ExitConversation, RequestsGet, WineDB from core.tools.editor import CodeEditor -from core.tools.cpu import ( - RequestsGet, - WineDB, - ExitConversation, -) from core.tools.gpu import ( ImageEditing, InstructPix2Pix, Text2Image, VisualQuestionAnswering, ) -from core.handlers.base import BaseHandler, FileHandler, FileType -from core.handlers.image import ImageCaptioning -from core.handlers.dataframe import CsvToDataframe +from core.tools.terminal import Terminal from core.upload import StaticUploader - +from env import settings from logger import logger app = FastAPI() diff --git a/core/agents/builder.py b/core/agents/builder.py index b4ad17e..14f43f2 100644 --- a/core/agents/builder.py +++ b/core/agents/builder.py @@ -1,14 +1,13 @@ from langchain.chat_models.base import BaseChatModel from langchain.output_parsers.base import BaseOutputParser -from env import settings - from core.prompts.input import EVAL_PREFIX, EVAL_SUFFIX from core.tools.base import BaseToolSet from core.tools.factory import ToolsFactory +from env import settings -from .llm import ChatOpenAI from .chat_agent import ConversationalChatAgent +from .llm import ChatOpenAI from .parser import EvalOutputParser diff --git a/core/agents/callback.py b/core/agents/callback.py index 387ced5..3f037d4 100644 --- a/core/agents/callback.py +++ b/core/agents/callback.py @@ -3,8 +3,8 @@ from typing import Any, Dict, List, Optional, Union from langchain.callbacks.base import BaseCallbackHandler from langchain.schema import AgentAction, AgentFinish, LLMResult -from logger import logger from ansi import ANSI, Color, Style +from logger import logger class EVALCallbackHandler(BaseCallbackHandler): diff --git a/core/agents/llm.py b/core/agents/llm.py index 7a4c6b5..27ab2f5 100644 --- a/core/agents/llm.py +++ b/core/agents/llm.py @@ -5,15 +5,6 @@ import logging import sys from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple -from pydantic import BaseModel, Extra, Field, root_validator -from tenacity import ( - before_sleep_log, - retry, - retry_if_exception_type, - stop_after_attempt, - wait_exponential, -) - from langchain.chat_models.base import BaseChatModel from langchain.schema import ( AIMessage, @@ -25,6 +16,14 @@ from langchain.schema import ( SystemMessage, ) from langchain.utils import get_from_dict_or_env +from pydantic import BaseModel, Extra, Field, root_validator +from tenacity import ( + before_sleep_log, + retry, + retry_if_exception_type, + stop_after_attempt, + wait_exponential, +) from logger import logger @@ -128,7 +127,7 @@ class ChatOpenAI(BaseChatModel, BaseModel): """Whether to stream the results or not.""" n: int = 1 """Number of chat completions to generate for each prompt.""" - max_tokens: int = 256 + max_tokens: int = 2048 """Maximum number of tokens to generate.""" class Config: @@ -226,7 +225,6 @@ class ChatOpenAI(BaseChatModel, BaseModel): def _generate( self, messages: List[BaseMessage], stop: Optional[List[str]] = None ) -> ChatResult: - message_dicts, params = self._create_message_dicts(messages, stop) logger.debug("Messages:\n") for item in message_dicts: diff --git a/core/agents/manager.py b/core/agents/manager.py index 52d500d..4f7ebf4 100644 --- a/core/agents/manager.py +++ b/core/agents/manager.py @@ -1,18 +1,17 @@ from typing import Dict -from langchain.agents.tools import BaseTool from langchain.agents.agent import Agent, AgentExecutor +from langchain.agents.tools import BaseTool +from langchain.callbacks import set_handler +from langchain.callbacks.base import CallbackManager from langchain.chains.conversation.memory import ConversationBufferMemory from langchain.memory.chat_memory import BaseChatMemory -from langchain.callbacks.base import CallbackManager -from langchain.callbacks import set_handler from core.tools.base import BaseToolSet from core.tools.factory import ToolsFactory -from .callback import EVALCallbackHandler from .builder import AgentBuilder - +from .callback import EVALCallbackHandler callback_manager = CallbackManager([EVALCallbackHandler()]) set_handler(EVALCallbackHandler()) diff --git a/core/handlers/base.py b/core/handlers/base.py index c4d906a..336b3d2 100644 --- a/core/handlers/base.py +++ b/core/handlers/base.py @@ -1,8 +1,9 @@ import os import uuid -import requests -from typing import Dict from enum import Enum +from typing import Dict + +import requests class FileType(Enum): diff --git a/core/handlers/image.py b/core/handlers/image.py index d4e8577..c507c05 100644 --- a/core/handlers/image.py +++ b/core/handlers/image.py @@ -1,9 +1,6 @@ import torch from PIL import Image -from transformers import ( - BlipProcessor, - BlipForConditionalGeneration, -) +from transformers import BlipForConditionalGeneration, BlipProcessor from core.prompts.file import IMAGE_PROMPT diff --git a/core/tools/base.py b/core/tools/base.py index 5468a16..07cec3d 100644 --- a/core/tools/base.py +++ b/core/tools/base.py @@ -1,8 +1,8 @@ -from typing import Callable, Tuple from enum import Enum +from typing import Callable, Tuple -from langchain.agents.tools import Tool, BaseTool from langchain.agents.agent import AgentExecutor +from langchain.agents.tools import BaseTool, Tool class ToolScope(Enum): diff --git a/core/tools/cpu.py b/core/tools/cpu.py index ab6e933..08a95ab 100644 --- a/core/tools/cpu.py +++ b/core/tools/cpu.py @@ -1,15 +1,13 @@ -from env import settings - import requests - -from llama_index.readers.database import DatabaseReader -from llama_index import GPTSimpleVectorIndex - from bs4 import BeautifulSoup +from llama_index import GPTSimpleVectorIndex +from llama_index.readers.database import DatabaseReader -from .base import tool, BaseToolSet, ToolScope, SessionGetter +from env import settings from logger import logger +from .base import BaseToolSet, SessionGetter, ToolScope, tool + class RequestsGet(BaseToolSet): @tool( diff --git a/core/tools/editor/__init__.py b/core/tools/editor/__init__.py index 2e4e5ce..e22241e 100644 --- a/core/tools/editor/__init__.py +++ b/core/tools/editor/__init__.py @@ -1,4 +1,4 @@ -from core.tools.base import tool, BaseToolSet +from core.tools.base import BaseToolSet, tool from logger import logger from .patch import CodePatcher diff --git a/core/tools/editor/patch.py b/core/tools/editor/patch.py index 08474e1..7bfd4dd 100644 --- a/core/tools/editor/patch.py +++ b/core/tools/editor/patch.py @@ -56,8 +56,12 @@ test.py|7,5|9,13|news_titles = [] test.py|11,16|11,16|_titles """ +import os +from pathlib import Path from typing import Tuple +from env import settings + class Position: separator = "," @@ -76,7 +80,7 @@ class PatchCommand: separator = "|" def __init__(self, filepath: str, start: Position, end: Position, content: str): - self.filepath: str = filepath + self.filepath: str = str(Path(settings["PLAYGROUND_DIR"]) / Path(filepath)) self.start: Position = start self.end: Position = end self.content: str = content @@ -92,6 +96,13 @@ class PatchCommand: return sum([len(line) for line in lines]) def execute(self) -> Tuple[int, int]: + # make sure the directory exists + if not str(Path(self.filepath).resolve()).startswith( + str(Path(settings["PLAYGROUND_DIR"]).resolve()) + ): + return "You can't write file outside of current directory." + + os.makedirs(os.path.dirname(self.filepath), exist_ok=True) lines = self.read_lines() before = sum([len(line) for line in lines]) diff --git a/core/tools/editor/read.py b/core/tools/editor/read.py index 42bac94..7b76cec 100644 --- a/core/tools/editor/read.py +++ b/core/tools/editor/read.py @@ -3,7 +3,7 @@ read protocol: |- """ -from typing import Optional, List, Tuple +from typing import List, Optional, Tuple class Line: diff --git a/core/tools/editor/write.py b/core/tools/editor/write.py index cf5fed9..f369afb 100644 --- a/core/tools/editor/write.py +++ b/core/tools/editor/write.py @@ -6,6 +6,7 @@ write protocol: """ import os from pathlib import Path + from env import settings diff --git a/core/tools/factory.py b/core/tools/factory.py index 21dbb3d..515f3ce 100644 --- a/core/tools/factory.py +++ b/core/tools/factory.py @@ -1,4 +1,5 @@ from typing import Optional + from langchain.agents import load_tools from langchain.agents.tools import BaseTool from langchain.llms.base import BaseLLM diff --git a/core/tools/gpu.py b/core/tools/gpu.py index 18cd79b..1a28468 100644 --- a/core/tools/gpu.py +++ b/core/tools/gpu.py @@ -1,29 +1,26 @@ import os import uuid + import numpy as np import torch -from PIL import Image - -from transformers import ( - CLIPSegProcessor, - CLIPSegForImageSegmentation, -) -from transformers import ( - BlipProcessor, - BlipForQuestionAnswering, -) - from diffusers import ( - StableDiffusionPipeline, + EulerAncestralDiscreteScheduler, StableDiffusionInpaintPipeline, StableDiffusionInstructPix2PixPipeline, + StableDiffusionPipeline, +) +from PIL import Image +from transformers import ( + BlipForQuestionAnswering, + BlipProcessor, + CLIPSegForImageSegmentation, + CLIPSegProcessor, ) -from diffusers import EulerAncestralDiscreteScheduler -from utils import get_new_image_name from logger import logger +from utils import get_new_image_name -from .base import tool, BaseToolSet +from .base import BaseToolSet, tool class MaskFormer(BaseToolSet): diff --git a/core/tools/terminal/__init__.py b/core/tools/terminal/__init__.py index 272e234..59d5b6d 100644 --- a/core/tools/terminal/__init__.py +++ b/core/tools/terminal/__init__.py @@ -1,12 +1,11 @@ import subprocess -from typing import Dict, List - from tempfile import TemporaryFile +from typing import Dict, List +from core.tools.base import BaseToolSet, SessionGetter, ToolScope, tool +from core.tools.terminal.syscall import SyscallTracer from env import settings from logger import logger -from core.tools.base import tool, BaseToolSet, ToolScope, SessionGetter -from core.tools.terminal.syscall import SyscallTracer class Terminal(BaseToolSet): @@ -17,9 +16,8 @@ class Terminal(BaseToolSet): name="Terminal", description="Executes commands in a terminal." "If linux errno occurs, we have to solve the problem with the terminal. " - "It can't execute interactive operations or blocking operations. " - "Input should be valid commands, " - "and the output will be any output from running that command.", + "Input must be one valid command. " + "Output will be any output from running that command.", scope=ToolScope.SESSION, ) def execute(self, commands: str, get_session: SessionGetter) -> str: @@ -45,9 +43,6 @@ class Terminal(BaseToolSet): except Exception as e: output = str(e) - if len(output) > 1000: - output = output[:1000] + "..." - logger.debug( f"\nProcessed Terminal, Input Commands: {commands} " f"Output Answer: {output}" diff --git a/core/tools/terminal/syscall.py b/core/tools/terminal/syscall.py index c38c27f..fd16399 100644 --- a/core/tools/terminal/syscall.py +++ b/core/tools/terminal/syscall.py @@ -1,16 +1,16 @@ -from typing import Tuple, Optional import signal +from typing import Optional, Tuple from ptrace.debugger import ( - PtraceDebugger, - PtraceProcess, - ProcessExit, - ProcessSignal, NewProcessEvent, ProcessExecution, + ProcessExit, + ProcessSignal, + PtraceDebugger, + PtraceProcess, ) -from ptrace.syscall import PtraceSyscall from ptrace.func_call import FunctionCallOptions +from ptrace.syscall import PtraceSyscall from ptrace.tools import signal_to_exitcode @@ -61,7 +61,7 @@ class SyscallTracer: break try: - self.wait_syscall_with_timeout(5) + self.wait_syscall_with_timeout(60) except ProcessExit as event: if event.exitcode is not None: exitcode = event.exitcode diff --git a/core/upload/s3.py b/core/upload/s3.py index 98f3335..9225d8d 100644 --- a/core/upload/s3.py +++ b/core/upload/s3.py @@ -1,7 +1,9 @@ import os + import boto3 from env import DotEnv + from .base import AbstractUploader diff --git a/core/upload/static.py b/core/upload/static.py index 414d517..1f35414 100644 --- a/core/upload/static.py +++ b/core/upload/static.py @@ -2,6 +2,7 @@ import os import shutil from env import DotEnv + from .base import AbstractUploader diff --git a/docker-compose.yml b/docker-compose.yml index f1e7365..4f4b93a 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -10,7 +10,7 @@ services: volumes: - ./static/:/app/static/ ports: - - "3000:3000" + - "4500:4500" - "7000:7000" - "8000:8000" # eval port env_file: @@ -26,7 +26,7 @@ services: - ../.cache/huggingface/:/root/.cache/huggingface/ - ./static/:/app/static/ ports: - - "3000:3000" + - "4500:4500" - "7000:7000" - "8000:8000" # eval port env_file: diff --git a/logger.py b/logger.py index f080525..349c2e9 100644 --- a/logger.py +++ b/logger.py @@ -1,4 +1,5 @@ import logging + from env import settings logger = logging.getLogger() diff --git a/utils.py b/utils.py index b3bc08f..e9c44df 100644 --- a/utils.py +++ b/utils.py @@ -1,9 +1,9 @@ import os import random -import torch import uuid -import numpy as np +import numpy as np +import torch os.makedirs("image", exist_ok=True) os.makedirs("audio", exist_ok=True)