from langchain_huggingface import HuggingFacePipeline as HF, ChatHuggingFace as Ch
from subprocess import Popen, PIPE as P
from langchain_experimental.tools.python.tool import PythonREPLTool as PYT
from langchain.agents import load_tools, create_structured_chat_agent as Agent,AgentExecutor as Ex, AgentType as Type
from langchain.agents.agent_toolkits import create_retriever_tool as crt
from langchain_community.agent_toolkits import FileManagementToolkit as FMT
from langchain.tools import Tool
from langchain.memory import ConversationBufferMemory as MEM,RedisChatMessageHistory as HIS
from langchain.schema import SystemMessage as SM,HumanMessage as HM, AIMessage as AM
from langchain import hub
import os
from langchain_core.prompts.chat import ChatPromptTemplate, MessagesPlaceholder
system = '''Respond to the human as helpfully and accurately as possible. You have access to the following tools:

{tools}

Use a json blob to specify a tool by providing an action key (tool name) and an action_input key (tool input).

Valid "action" values: "Final Answer" or {tool_names}

Provide only ONE action per $JSON_BLOB, as shown:

```
{{
  "action": $TOOL_NAME,
  "action_input": $INPUT
}}
```

Follow this format:

Question: input question to answer
Thought: consider previous and subsequent steps
Action:
```
$JSON_BLOB
```
Observation: action result
... (repeat Thought/Action/Observation N times)
Thought: I know what to respond
Action:
```
{{
  "action": "Final Answer",
  "action_input": "Final response to human"
}}

Begin! Reminder to ALWAYS respond with a valid json blob of a single action. Use tools if necessary. Respond directly if appropriate. Format is Action:```$JSON_BLOB```then Observation'''

human = '''

{input}

{agent_scratchpad}

(reminder to respond in a JSON blob no matter what)'''

prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system),
        MessagesPlaceholder("chat_history", optional=True),
        ("human", human),
    ]
)
from typing import Any, Dict, List, Optional
from langchain_core.language_models import BaseChatModel
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage
from langchain_core.outputs import ChatResult, ChatGeneration
from langchain_core.callbacks.manager import CallbackManagerForLLMRun
from langchain_core.callbacks.manager import AsyncCallbackManagerForLLMRun
from langchain_core.runnables import run_in_executor
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
class Chatchat(BaseChatModel):

    model_name: str = "peterpeter8585/deepseek_1"
    tokenizer : AutoTokenizer = None
    model: AutoModelForCausalLM = None
    model_path: str = None

    def __init__(self, model_path, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        if model_path is not None:
            self.model_name = model_path

        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            self.model_name, trust_remote_code=True)


    def _call(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> str:
        # Load and preprocess the image
        messages = [
            {"role": "system", "content": "You are Chatchat.A helpful assistant at code."},
            {"role": "user", "content": prompt}
        ]

        text = self.tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True
        )
        model_inputs  = self.tokenizer([text], return_tensors="pt").to(self.model.device)
        generated_ids = self.model.generate(
              **model_inputs,
              max_new_tokens=512
          )
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
          ]

        response = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

        return response

    async def _acall(
            self,
            prompt: str,
            stop: Optional[List[str]] = None,
            run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> str:
        # Implement the async logic to generate a response from the model
        return await run_in_executor(
            None,
            self._call,
            prompt,
            stop,
            run_manager.get_sync() if run_manager else None,
            **kwargs,
        )

    @property
    def _llm_type(self) -> str:
        return "custom-llm-chat"

    @property
    def _identifying_params(self) -> Dict[str, Any]:
        return {"model_name": self.model_name}

    def _generate(
            self,
            messages: List[BaseMessage],
            stop: Optional[List[str]] = None,
            run_manager: Optional[CallbackManagerForLLMRun] = None,
            **kwargs: Any,
    ) -> ChatResult:
        # Assumes the first message contains the prompt and the image path is in metadata
        prompt = messages[0].content
        response_text = self._call(prompt,  stop, run_manager, **kwargs)

        # Create AIMessage with the response
        ai_message = AIMessage(content=response_text)
        return ChatResult(generations=[ChatGeneration(message=ai_message)])


llm=Chatchat(model_path=None)
#from transformers import pipeline,AutoModelForCausalLM as M,AutoTokenizer as T
#m=M.from_pretrained("peterpeter8585/syai4.3")
#t=T.from_pretrained("peterpeter8585/syai4.3")
#pipe=pipeline(model=m,tokenizer=t,task="text-generation")
#llm=HF.from_model_id(model_id="peterpeter8585/syai4.6",task="text-generation")
from langchain.retrievers import WikipediaRetriever as Wiki
import gradio as gr
chatbot = gr.Chatbot(
    label="SYAI4.1",
    show_copy_button=True,
    layout="panel"
)
def terminal(c):
    a=Popen(c,shell=True,stdin=P,stdout=P,stderr=P)
    return a.stdout.read()+a.stderr.read()
tools=FMT().get_tools()
tools.append(PYT())
tools.extend(load_tools(["requests_all"],allow_dangerous_tools=True))
tools.extend(load_tools(["llm-math","ddg-search"],llm=llm))
tools.append(Tool.from_function(func=terminal,name="terminal",description="터미널 명령어실행에 적합함"))
tools.append(crt(name="wiki",description="위키 백과를 검색하여 정보를 가져온다",retriever=Wiki(lang="ko",top_k_results=1)))
def chat(message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p, chat_session):
    messages=[SM(content=system_message+"And, Your name is Chatchat")]
    for val in history:
        if val[0]:
            messages.append(HM(content=val[0]))
        if val[1]:
            messages.append(AM(content=val[1]))

    messages.append(HM(content=message))
    memory=MEM(memory_key="history")
    agent=Ex(agent=Agent(llm,tools,prompt),tools=tools,verbose=True,handle_parsing_errors=True,memory=memory)
    return agent.invoke({"input":messages,"chat_history":memory.buffer_as_messages})
ai1=gr.ChatInterface(
    chat,
    chatbot=chatbot,
    additional_inputs=[
        gr.Textbox(value="You are a helpful assistant.", label="System message",  interactive=True),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.1, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.1,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
        gr.Textbox(label="chat_id(please enter the chat id!)")
    ],
    
)
ai1.launch()