feat: adds hf model

tdagent/grchat.py

@@ -10,6 +10,7 @@ import botocore.exceptions
 import gradio as gr
 from langchain_aws import ChatBedrock
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from langchain_huggingface import HuggingFaceEndpoint
 from langchain_mcp_adapters.client import MultiServerMCPClient
 from langgraph.prebuilt import create_react_agent
 
@@ -56,6 +57,7 @@ llm_agent: CompiledGraph | None = None
 #### Utility functions ####
 
 
+## Bedrock LLM creation ##
 def create_bedrock_llm(
     bedrock_model_id: str,
     aws_access_key: str,
@@ -91,9 +93,25 @@ def create_bedrock_llm(
     return llm, ""
 
 
-#### UI functionality ####
+## Hugging Face LLM creation ##
+def create_hf_llm(
+    hf_model_id: str,
+    huggingfacehub_api_token: str | None = None,
+) -> tuple[HuggingFaceEndpoint | None, str]:
+    """Create a LangGraph Hugging Face agent."""
+    try:
+        llm = HuggingFaceEndpoint(
+            model=hf_model_id,
+            huggingfacehub_api_token=huggingfacehub_api_token,
+            temperature=0.8,
+        )
+    except Exception as e:  # noqa: BLE001
+        return None, str(e)
+
+    return llm, ""
 
 
+#### UI functionality ####
 async def gr_connect_to_bedrock(
     model_id: str,
     access_key: str,
@@ -128,7 +146,7 @@ async def gr_connect_to_bedrock(
     #     }
     # )
     # tools = await client.get_tools()
-
+    tools = []
     if mcp_servers:
         client = MultiServerMCPClient(
             {
@@ -140,8 +158,6 @@ async def gr_connect_to_bedrock(
             },
         )
         tools = await client.get_tools()
-    else:
-        tools = []
 
     llm_agent = create_react_agent(
         model=llm,
@@ -152,6 +168,40 @@ async def gr_connect_to_bedrock(
     return "✅ Successfully connected to AWS Bedrock!"
 
 
+async def gr_connect_to_hf(
+    model_id: str,
+    hf_access_token_textbox: str | None,
+    mcp_servers: Sequence[MutableCheckBoxGroupEntry] | None,
+) -> str:
+    """Initialize Hugging Face agent."""
+    global llm_agent  # noqa: PLW0603
+
+    llm, error = create_hf_llm(model_id, hf_access_token_textbox)
+
+    if llm is None:
+        return f"❌ Connection failed: {error}"
+    tools = []
+    if mcp_servers:
+        client = MultiServerMCPClient(
+            {
+                server.name.replace(" ", "-"): {
+                    "url": server.value,
+                    "transport": "sse",
+                }
+                for server in mcp_servers
+            },
+        )
+        tools = await client.get_tools()
+
+    llm_agent = create_react_agent(
+        model=llm,
+        tools=tools,
+        prompt=SYSTEM_MESSAGE,
+    )
+
+    return "✅ Successfully connected to Hugging Face!"
+
+
 async def gr_chat_function(  # noqa: D103
     message: str,
     history: list[Mapping[str, str]],
@@ -249,6 +299,32 @@ with gr.Blocks() as gr_app:
             outputs=[status_textbox],
         )
 
+    with gr.Accordion("Hugging Face Configuration", open=True):
+        with gr.Row():
+            hf_model_id_textbox = gr.Textbox(
+                label="HF Model Id",
+                value="fdtn-ai/Foundation-Sec-8B",
+            )
+        with gr.Row():
+            hf_access_token_textbox = gr.Textbox(
+                label="Hugging Face Access Token",
+                type="password",
+                placeholder="Enter your Hugging Face Access Token",
+            )
+            hf_connect_btn = gr.Button("🔌 Connect to Hugging Face", variant="primary")
+
+        status_textbox = gr.Textbox(label="Connection Status", interactive=False)
+
+        hf_connect_btn.click(
+            gr_connect_to_hf,
+            inputs=[
+                hf_model_id_textbox,
+                hf_access_token_textbox,
+                mcp_list.state,
+            ],
+            outputs=[status_textbox],
+        )
+
     chat_interface = gr.ChatInterface(
         fn=gr_chat_function,
         type="messages",