[CM-10413] Fix API.query type hints to the query parameter (#142)

* [CM-10413] Fix API.query type hints to the query parameter

* [CM-10413] change import to avoid pyright linter error

src/comet_llm/api_objects/api.py

@@ -7,14 +7,15 @@
 #  \____\___/|_| |_| |_|\___|\__(_)_| |_| |_|_|
 #
 #  Sign up for free at https://www.comet.com
-#  Copyright (C) 2015-2023 Comet ML INC
+#  Copyright (C) 2015-2024 Comet ML INC
 #  This source code is licensed under the MIT license found in the
 #  LICENSE file in the root directory of this package.
 # *******************************************************
 
 from typing import List, Optional
 
 import comet_ml
+from comet_ml.query import QueryExpression
 
 from .. import experiment_info, logging_messages, query_dsl
 from . import llm_trace_api
@@ -28,7 +29,7 @@ def __init__(self, api_key: Optional[str] = None) -> None:
         API class for accessing and updating prompt information.
 
         Args:
-            api_key: You private Comet API key
+            api_key: Your private Comet API key
 
         Example:
             ```python linenums="1"
@@ -91,29 +92,28 @@ def get_llm_trace_by_name(
         return llm_trace_api.LLMTraceAPI.__api__from_api_experiment__(matching_trace[0])
 
     def query(
-        self, workspace: str, project_name: str, query: str
+        self, workspace: str, project_name: str, query: QueryExpression
     ) -> List[llm_trace_api.LLMTraceAPI]:
         """
-        Fetch LLM Trace based on a query. Currently it is only possible to use
+        Fetch LLM Trace based on a query. Currently, it is only possible to use
         trace metadata or details fields to filter the traces.
 
         Args:
             workspace: Name of the workspace
             project_name: Name of the project
-            query: Query to use
+            query: Query expression to use (see note below)
 
         Returns: A list of LLMTrace objects
 
         Note:
-            The `query` object takes the form of (QUERY_VARIABLE OPERATOR VALUE) with:
+            The `query` object can also take the form of (QUERY_VARIABLE OPERATOR VALUE) with:
 
             * QUERY_VARIABLE is either TraceMetadata, Duration, Timestamp.
             * OPERATOR is any standard mathematical operators `<=`, `>=`, `!=`, `<`, `>`.
 
             It is also possible to add multiple query conditions using `&`.
 
-            If you are querying nested parameters, you should flatted the parameter name using the
-            `.` operator.
+            If you are querying nested parameters, flatten the parameter name using the `.` operator.
 
             To query the duration, you can use Duration().