optimize: replace json.loads with TypeAdapter for JSON unmarshaling (#103)

* refactor: replace json.loads with TypeAdapter for JSON validation in request readers and model

- Updated JSON parsing in BackwardsInvocation, StdioRequestReader, and TCPReaderWriter to use TypeAdapter for improved validation.
- Enhanced error handling in OAICompatLargeLanguageModel by replacing json.loads with TypeAdapter, ensuring better compliance with expected JSON schema.

* fix: linter

Author: Yeuoly

python/dify_plugin/core/runtime.py

@@ -1,13 +1,12 @@
-import json
 import uuid
 from abc import ABC
 from collections.abc import Generator, Mapping
 from concurrent.futures import ThreadPoolExecutor
 from enum import Enum
-from typing import Generic, Optional, TypeVar, Union
+from typing import Any, Generic, Optional, TypeVar, Union
 
 import httpx
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from yarl import URL
 
 from dify_plugin.config.config import InstallMethod
@@ -276,7 +275,7 @@ def generator():
                     if not line:
                         continue
 
-                    data = json.loads(line)
+                    data = TypeAdapter(dict[str, Any]).validate_json(line)
                     yield PluginInStreamBase(
                         session_id=data["session_id"],
                         event=PluginInStreamEvent.value_of(data["event"]),

python/dify_plugin/core/server/stdio/request_reader.py

@@ -1,8 +1,9 @@
 import sys
 from collections.abc import Generator
-from json import loads
+from typing import Any
 
 from gevent.os import tp_read
+from pydantic import TypeAdapter
 
 from dify_plugin.core.entities.plugin.io import (
     PluginInStream,
@@ -37,7 +38,7 @@ def _read_stream(self) -> Generator[PluginInStream, None, None]:
             lines = lines[:-1]
             for line in lines:
                 try:
-                    data = loads(line)
+                    data = TypeAdapter(dict[str, Any]).validate_json(line)
                     yield PluginInStream(
                         session_id=data["session_id"],
                         event=PluginInStreamEvent.value_of(data["event"]),

python/dify_plugin/core/server/tcp/request_reader.py

@@ -5,13 +5,13 @@
 import socket as native_socket
 import time
 from collections.abc import Callable, Generator
-from json import loads
 from threading import Lock
-from typing import Optional
+from typing import Any, Optional
 
 from gevent import sleep
 from gevent import socket as gevent_socket
 from gevent.select import select
+from pydantic import TypeAdapter
 
 from dify_plugin.core.entities.message import InitializeMessage
 from dify_plugin.core.entities.plugin.io import (
@@ -190,7 +190,7 @@ def _read_stream(self) -> Generator[PluginInStream, None, None]:
             lines = lines[:-1]
             for line in lines:
                 try:
-                    data = loads(line)
+                    data = TypeAdapter(dict[str, Any]).validate_json(line)
                     chunk = PluginInStream(
                         session_id=data["session_id"],
                         event=PluginInStreamEvent.value_of(data["event"]),

python/dify_plugin/interfaces/model/openai_compatible/llm.py

@@ -3,10 +3,11 @@
 import logging
 from collections.abc import Generator
 from decimal import Decimal
-from typing import Optional, Union, cast
+from typing import Any, Optional, Union, cast
 from urllib.parse import urljoin
 
 import requests
+from pydantic import TypeAdapter, ValidationError
 
 from dify_plugin.entities import I18nObject
 from dify_plugin.entities.model import (
@@ -356,7 +357,7 @@ def _generate(
                 if not json_schema:
                     raise ValueError("Must define JSON Schema when the response format is json_schema")
                 try:
-                    schema = json.loads(json_schema)
+                    schema = TypeAdapter(dict[str, Any]).validate_json(json_schema)
                 except Exception as exc:
                     raise ValueError(f"not correct json_schema format: {json_schema}") from exc
                 model_parameters.pop("json_schema")
@@ -517,9 +518,9 @@ def _handle_generate_stream_response(
                     continue
 
                 try:
-                    chunk_json: dict = json.loads(decoded_chunk)
+                    chunk_json: dict = TypeAdapter(dict[str, Any]).validate_json(decoded_chunk)
                 # stream ended
-                except json.JSONDecodeError:
+                except ValidationError:
                     yield self._create_final_llm_result_chunk(
                         index=chunk_index + 1,
                         message=AssistantPromptMessage(content=""),