app.py

# app_pixal_chat.py
import os, re, json, gradio as gr, requests
from typing import Optional, List, Dict, Any
from requests.adapters import HTTPAdapter, Retry
from langchain.llms.base import LLM
from langchain.memory import ConversationBufferMemory
from langchain.agents import initialize_agent, AgentType, load_tools
from langchain.tools import Tool
from langchain_experimental.tools.python.tool import PythonREPLTool
from langchain_community.retrievers import WikipediaRetriever

# ──────────────────────────────
# ✅ GitHubModelLLM (그대로 유지)
# ──────────────────────────────

from typing import Optional, List, Dict, Any
from langchain.llms.base import LLM
import requests, os, json
from requests.adapters import HTTPAdapter, Retry

class GitHubModelLLM(LLM):
    """GitHub Models API 기반 LangChain LLM (대화 메모리 통합 지원)"""
    model: str = "openai/gpt-4.1"
    endpoint: str = "https://models.github.ai/inference"
    token: Optional[str] = os.environ.get("token")
    system_prompt: Optional[str] = "너는 PIXAL(Primary Interactive X-ternal Assistant with multi Language)이야. 너의 개발자는 정성윤 이라는 6학년 파이썬 프로그래머야.이것은 시스템 메시지입니다.참고 하십시오."
    request_timeout: float = 30.0
    max_retries: int = 2
    backoff_factor: float = 0.3

    @property
    def _llm_type(self) -> str:
        return "github_models_api"

    def _post_chat(self, body: Dict[str, Any]) -> Dict[str, Any]:
        token = self.token or os.getenv("GITHUB_TOKEN") or os.getenv("token")
        session = requests.Session()
        retries = Retry(total=self.max_retries, backoff_factor=self.backoff_factor,
                        status_forcelist=[429, 500, 502, 503, 504])
        session.mount("https://", HTTPAdapter(max_retries=retries))
        session.headers.update({
            "Content-Type": "application/json",
            "Authorization": f"Bearer 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"
        })
        resp = session.post(f"{self.endpoint}/chat/completions", json=body, timeout=self.request_timeout)
        resp.raise_for_status()
        return resp.json()

    def _call(
        self,
        prompt: str,
        stop: Optional[List[str]] = None,
        **kwargs
    ) -> str:
        """대화 메모리(chat_history)를 포함하여 모델 호출"""
        # 💬 메모리에 저장된 대화 메시지 불러오기
        memory = kwargs.get("memory")
        messages = []

        if self.system_prompt:
            messages.append({"role": "system", "content": self.system_prompt})

        # memory가 있을 경우 (이전 대화 포함)
        if memory and hasattr(memory, "chat_memory"):
            for msg in memory.chat_memory.messages:
                role = "user" if msg.type == "human" else "assistant"
                messages.append({"role": role, "content": msg.content})

        # 현재 사용자 입력
        messages.append({"role": "user", "content": prompt})

        body = {"model": self.model, "messages": messages}
        if stop:
            body["stop"] = stop

        # API 호출
        res = self._post_chat(body)
        msg = res.get("choices", [{}])[0].get("message", {})
        return msg.get("content") or json.dumps(msg.get("function_call", {}))
# ──────────────────────────────
# ✅ LangChain 도구 & 에이전트 구성
# ──────────────────────────────
llm = GitHubModelLLM()
tools = load_tools(["ddg-search", "requests_all", "llm-math"], llm=llm,allow_dangerous_tools=True)
tools.append(Tool(name="python_repl", func=PythonREPLTool().run, description="Python 코드 실행 도구"))
retriever = WikipediaRetriever(lang="ko")
tools.append(Tool(name="wiki", func=retriever.get_relavant_documents, description="위키백과 검색"))

# ✅ 대화 기억 메모리
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# ✅ Agent (Memory 연동)
agent = initialize_agent(
    tools,
    llm,
    agent_type=AgentType.STRUCTURED_CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    memory=memory,
    verbose=True
)

# ──────────────────────────────
# ✅ Chat 함수 (Memory 유지)
# ──────────────────────────────
def chat(message, history):
    try:
        # 대화 기록을 LangChain memory에 반영
        memory.chat_memory.add_user_message(message)
        raw = agent.run(message)
        # JSON 형태로 반환 시 파싱
        text = str(raw)
        match = re.search(r"\{.*\}", text, re.DOTALL)
        if match:
            try:
                obj = json.loads(match.group(0))
                text = obj.get("action_input") or obj.get("Final Answer") or obj.get("content") or text
            except Exception:
                pass
        # AI 응답을 memory에 추가
        memory.chat_memory.add_ai_message(text)
    except Exception as e:
        text = f"⚠️ 오류: {e}"

    history = history + [(message, text)]
    return history, history, ""

# ──────────────────────────────
# ✅ Gradio UI (ChatGPT 스타일)
# ──────────────────────────────
with gr.Blocks(theme=gr.themes.Soft(), title="PIXAL Assistant") as demo:
    gr.HTML("""
    <div style="background:#f1f5f9;padding:12px;border-bottom:1px solid #d1d5db;
                display:flex;align-items:center;justify-content:space-between;">
        <h2 style="margin:0;">🤖 PIXAL Assistant</h2>
        <span style="font-size:0.9em;color:#555;">LangChain + GitHub LLM</span>
    </div>
    """)

    chatbot = gr.Chatbot(
        label=None,
        height=720,
        bubble_full_width=False,
        render_markdown=True,
        avatar_images=("https://avatars.githubusercontent.com/u/9919?s=280&v=4", None),
    )

    with gr.Row():
        msg = gr.Textbox(placeholder="메시지를 입력하세요...", show_label=False, scale=8)
        send = gr.Button("전송", variant="primary", scale=1)
        clear = gr.Button("🧹 초기화", scale=1)

    msg.submit(chat, [msg, chatbot], [chatbot, chatbot, msg])
    send.click(chat, [msg, chatbot], [chatbot, chatbot, msg])
    clear.click(lambda: None, None, chatbot, queue=False)

    gr.Markdown("""
    <div style="text-align:center;color:#777;font-size:0.85em;margin-top:8px;">
    💡 대화 기록은 세션 동안 유지됩니다.  
    Made with ❤️ by PIXAL
    </div>
    """)

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)